From 649ba8032be8719346b4f4e30bc5d2fb1ef33591 Mon Sep 17 00:00:00 2001 From: IceRhymers Date: Sat, 25 Jul 2026 01:14:49 -0700 Subject: [PATCH 1/9] chore: initialize indexer performance integration branch From 86e64d792a32f81b589574ab19e2c189e2715229 Mon Sep 17 00:00:00 2001 From: Tanner <84605639+IceRhymers@users.noreply.github.com> Date: Sat, 25 Jul 2026 07:03:13 -0700 Subject: [PATCH 2/9] indexer: per-phase timing instrumentation (#103) Refs #103 --- docs/runbooks/indexing-parallelism.md | 58 +- indexer/AGENTS.md | 7 +- indexer/job.py | 179 ++++++- indexer/store.py | 10 + indexer/timing.py | 117 +++++ tests/integration/test_store.py | 33 ++ tests/unit/test_job.py | 730 ++++++++++++++++++++++++++ tests/unit/test_timing.py | 86 +++ 8 files changed, 1211 insertions(+), 9 deletions(-) create mode 100644 indexer/timing.py create mode 100644 tests/unit/test_timing.py diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index f4e9947..fed089c 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/docs/runbooks/indexing-parallelism.md @@ -116,8 +116,9 @@ loop, third-party libraries) carry `-`. ``` INFO indexer.job [-]: local disk at /tmp: 41.2 GB free of 64.0 GB total; 4 worker(s) x 2.5 GB peak INFO indexer.fetch [acme/widgets]: ... -INFO indexer.job [acme/widgets]: finished acme/widgets in 71.30s -INFO indexer.job [acme/gadgets]: skipped acme/gadgets: already indexed at abc123 (semantics v1) in 0.41s +INFO indexer.job [acme/widgets]: phase timing acme/widgets@main: total=213.32s resolve=0.00s download=12.10s extract=8.40s parse=31.00s embed=88.20s db=64.50s sweep=0.30s other=8.82s +INFO indexer.job [acme/widgets]: finished acme/widgets in 213.74s (resolve=0.42s list=0.00s) +INFO indexer.job [acme/gadgets]: skipped acme/gadgets@main: already indexed at abc123 (semantics v1) in 0.41s ``` **To find the giant:** grep for `finished .* in` and sort by the elapsed number. @@ -130,6 +131,59 @@ accept the duration. line against the sum of the per-repo elapsed times. If the total is already close to the slowest single repo, the pool is not the bottleneck. +### 2.1 Finding the dominant phase, not just the dominant repo + +Every **indexed** branch emits one `phase timing` line accounting for its entire +wall clock. Skipped, failed, and conflicted branches emit none — there is nothing +to attribute. + +``` +grep 'phase timing' run.log # one line per indexed branch +``` + +The nine fields are fixed, always present, always in this order, always `%.2fs`. +A phase that did not run prints `0.00s` rather than disappearing, so the line +never changes shape between a semantic-on and a semantic-off run and every grep +you write keeps working. Read the largest field; that is the branch's bottleneck. + +| Dominant phase | What it means | Which issue addresses it | +|---|---|---| +| `download` / `extract` | archive I/O bound | #106 (single-pass in-memory ingestion) | +| `parse` | GIL-bound tree-sitter extraction | #108 (process-pool extraction) | +| `embed` | serial AI Gateway round trips | #107 (concurrent embedding) | +| `db` | per-file round trips | #105 (batched writes) | +| any of the above, on **unchanged** content | redundant work | #104 (file-level delta indexing) | + +Four fields need interpretation before you act on them: + +- **`resolve=0.00s` on a default branch is expected, not a bug.** That branch's + HEAD SHA came from the repo-level resolve, which happens once per repo outside + every branch's total and is reported on the repo's `finished` line as + `resolve=`. The `list=` on the same line is the branch-listing API call, which + is `0.00s` unless the repo has `branches:` globs configured (it is not called + at all otherwise) and which is paginated — on a monorepo with hundreds of + branches it is a real, and otherwise invisible, cost. The elapsed value in + `finished … in Xs` is measured on its own clock and is deliberately not + reconciled against the parenthesised numbers. +- **`other=` is the unattributed residual**, `total` minus every measured phase, + clamped at zero. It is dominated by the temp-dir teardown — an `rm -rf` of a + freshly extracted multi-GB tree — plus the pre-flight disk check. It exists so + the line has no silently missing time; a large `other` means something real is + happening outside every instrumented phase and is worth chasing. +- **`embed=` covers chunking as well as the network.** It spans `iter_chunks` + (CPU/GIL-bound) *and* the serial AI Gateway round trips. #107 addresses only + the round trips, so before routing work there, confirm the phase is + network-bound rather than chunking-bound (a follow-up may split it into + `chunk=`/`embed=`). +- **`db=` excludes parse and sweep, but the walk still happens inside the + transaction.** Files stream lazily through `index_repo`'s open transaction for + bounded memory, so file production is timed separately and subtracted from + `db`; the sweep is subtracted too. On the **non-semantic** path, though, the + directory walk itself materializes inside that open transaction (`parse.py`'s + `rglob`, on the first item). That is long-standing behavior which this + instrumentation merely makes visible for the first time — it is not a new + regression. + --- ## 3. The three limits, and why raising concurrency is a bad trade diff --git a/indexer/AGENTS.md b/indexer/AGENTS.md index eb5c590..f894434 100644 --- a/indexer/AGENTS.md +++ b/indexer/AGENTS.md @@ -14,13 +14,14 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` | `chunk_store.py` | `write_chunks`: delete-and-reinsert one file's rows in the `chunks` table (no natural key). Takes a live `Connection`; never opens an engine, never calls the embedder — vectors arrive precomputed. No `repo_id` parameter: `chunks` is scoped by `file_id` only. | | `fetch.py` | All GitHub HTTP: paginated org/user repo enumeration (`RepoMeta`), branch listing, `resolve_ref`/`resolve_branch_head` (branch name -> immutable SHA), streamed tarball download capped at `MAX_TARBALL_BYTES` (500 MB), safe extraction (`filter="data"`, bomb check) capped at `MAX_EXTRACTED_BYTES` (2 GB), `assert_disk_headroom` (2.5 GB per worker, both caps alive at once). `RateLimitError` is deliberately narrow: 429 always; 403 only with `Retry-After` or `X-RateLimit-Remaining: 0` — other 403s are permission failures. | | `hashing.py` | `content_sha`: canonical SHA-256 hex of content (`None` -> empty string). Single source of truth for `files.content_sha`; must stay byte-identical to the `0003` migration's SQL backfill forever or cross-branch dedup silently breaks (`tests/integration/test_content_sha_parity.py` is the gate). | -| `job.py` | Entry point and orchestration: `main()`/`run()`, `ThreadPoolExecutor` sized by `effective_workers` (unit of work = one repo, all branches sequential), batched `repo_branches` stamp read, per-branch skip-if-unchanged, per-branch `BranchOutcome` classification (`indexed`/`skipped`/`conflict`/`failed`), semantic precompute (`_precompute_chunk_writer`), `ContextVar`-based `[repo]` log attribution. Returns 1 if any branch failed; conflicts self-heal and do not fail the run. **Reconciliation checkpoint**: after every worker joins, `_decide_reconciliation` gates on zero failures/conflicts, every repo accounted for, and every repo's `discovery_complete`; on a pass, `_reconcile` reuses the pre-fan-out stamp snapshot to compute each repo's retired branches and the run's full desired repo set, then calls the injected `reconcile_retired_fn`/`reconcile_removed_fn` seams (default: the real `indexer.store` primitives) on one post-fan-out connection — retired branches first, repo purge second. `MAX_PURGE_SHRINK_FRACTION = 0.5` (hardcoded, no config knob) withholds ONLY the purge, as a logged incident signal, if it would remove a strict majority of currently stored repos; retired-branch cleanup on survivors still applies. `ReconcileProgress.committed_any` drives an honest "partially reconciled" vs "left stale" failure message — never the raw exception, only its phase/repo/type name. Never called from worker-thread code (`_index_one*`) — a source-level tripwire test enforces it. | +| `job.py` | Entry point and orchestration: `main()`/`run()`, `ThreadPoolExecutor` sized by `effective_workers` (unit of work = one repo, all branches sequential), batched `repo_branches` stamp read, per-branch skip-if-unchanged, per-branch `BranchOutcome` classification (`indexed`/`skipped`/`conflict`/`failed`), semantic precompute (`_precompute_chunk_writer`), `ContextVar`-based `[repo]` log attribution, and per-phase timing instrumentation (one `phase timing repo@branch: total=… resolve=… download=… extract=… parse=… embed=… db=… sweep=… other=…` INFO line per INDEXED branch, emitted from inside the worker so `[repo]` resolves; fixed unconditional field set; `_timed_items` charges lazy item production to `parse` so it is not fused into `db`; `resolve=`/`list=` on the per-repo `finished` line carry the repo-scoped costs that belong to no branch's total). Returns 1 if any branch failed; conflicts self-heal and do not fail the run. **Reconciliation checkpoint**: after every worker joins, `_decide_reconciliation` gates on zero failures/conflicts, every repo accounted for, and every repo's `discovery_complete`; on a pass, `_reconcile` reuses the pre-fan-out stamp snapshot to compute each repo's retired branches and the run's full desired repo set, then calls the injected `reconcile_retired_fn`/`reconcile_removed_fn` seams (default: the real `indexer.store` primitives) on one post-fan-out connection — retired branches first, repo purge second. `MAX_PURGE_SHRINK_FRACTION = 0.5` (hardcoded, no config knob) withholds ONLY the purge, as a logged incident signal, if it would remove a strict majority of currently stored repos; retired-branch cleanup on survivors still applies. `ReconcileProgress.committed_any` drives an honest "partially reconciled" vs "left stale" failure message — never the raw exception, only its phase/repo/type name. Never called from worker-thread code (`_index_one*`) — a source-level tripwire test enforces it. | | `languages.py` | Single source of truth shared by `parse.py` and `symbols.py`: `EXT_TO_LANG` (8 extensions -> tree-sitter language names), `SYMBOL_KINDS` (node type -> symbol kind per language), `EDGE_NODE_KINDS` (node type -> reference-edge kind per language, Python-only until #85), `MAX_FILE_BYTES` (1 MB), `SEMANTIC_CHUNK_MAX_CHARS` (2000, ~4 chars/token), and the frozen dataclasses `ParsedFile`, `Chunk`, `ExtractedSymbol`, `ExtractedEdge`, `FileExtraction`, `IndexCounts` (now carries `edges`). | | `parse.py` | `iter_source_files`: walk an extracted tree, yield every text file (unknown extensions kept with `lang=None`, since grep runs over all files); skips `.git/`, symlinks, files > `MAX_FILE_BYTES` (stat before read), NUL-sniffed binaries, and UTF-8 decode failures; strips surviving NULs (Postgres `text` rejects them). `iter_chunks`: deterministic line-aligned chunking, no overlap, no mid-line splits, 1-based inclusive line ranges. | | `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (2.5 GB peak per worker), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. Deliberately import-light: pydantic + PyYAML + stdlib only. | | `resolve.py` | `resolve_repos`: enumerate org/user selectors, apply `ExcludeRules` to enumerated repos only (explicit `repos` entries always win, unfiltered), dedup case-insensitively keeping first-seen spelling, union each connection's `branches:` globs into per-repo `RepoEntry.branch_globs`, then fail fast: `EmptyConfigError` on zero repos (indexing nothing must not exit 0), `RepoCeilingError` above `MAX_REPOS` (500, overridable via `--max_repos`). Also matches `RepoConfig.semantic_max_chunks_per_repo` case-insensitively onto each resolved repo's `RepoEntry.semantic_max_chunks` (`None` if unmatched); an override key matching no resolved repo logs a WARNING (typo guard). Enumerator params are a test seam only, not provider dispatch. | -| `store.py` | `index_repo`: the single atomic unit of work for one (repo, branch) inside `with conn.begin():` — repos upsert, `repo_branches` CAS-baseline read, per-file array-union upsert on `(repo_id, path, content_sha)` with branch-membership union, delete-and-reinsert `symbols` AND `reference_edges` (both keyed only by `file_id`, no natural key), optional `chunk_writer` call, membership sweep (strip this branch from unseen rows, delete rows with empty `branches`; skipped with a WARNING on an empty seen-set), then CAS stamp (raises `StaleIndexError` on mismatch, rolling everything back). `items` is `Iterable[tuple[ParsedFile, FileExtraction]]`; the `reference_edges` delete runs unconditionally, even when a file's `FileExtraction.edges` is empty, so stale rows never survive a re-index. Pure DML, no TEMP tables (job role has no TEMP privilege on Lakebase). `reconcile_retired_branches` and `reconcile_removed_repos`: the desired-state storage primitives `job.py`'s reconciliation checkpoint calls — each its own `conn.begin()`, repo-scoped membership subtraction (never delete-by-path) and an exact-match repo purge respectively; neither decides WHAT is retired/desired, only applies a caller-supplied set (see `job.py`'s row). | +| `store.py` | `index_repo`: the single atomic unit of work for one (repo, branch) inside `with conn.begin():` — repos upsert, `repo_branches` CAS-baseline read, per-file array-union upsert on `(repo_id, path, content_sha)` with branch-membership union, delete-and-reinsert `symbols` AND `reference_edges` (both keyed only by `file_id`, no natural key), optional `chunk_writer` call, membership sweep (strip this branch from unseen rows, delete rows with empty `branches`; skipped with a WARNING on an empty seen-set), then CAS stamp (raises `StaleIndexError` on mismatch, rolling everything back). `items` is `Iterable[tuple[ParsedFile, FileExtraction]]`; the `reference_edges` delete runs unconditionally, even when a file's `FileExtraction.edges` is empty, so stale rows never survive a re-index. The sweep call site is wall-clocked into `indexer.timing`'s ambient `PhaseTimer` (`record("sweep", …)`) — never into `IndexCounts` (frozen, compared by value) and never through a new `index_repo` parameter (that signature is an injected seam); `store.py` emits **no** log record for it. Pure DML, no TEMP tables (job role has no TEMP privilege on Lakebase). `reconcile_retired_branches` and `reconcile_removed_repos`: the desired-state storage primitives `job.py`'s reconciliation checkpoint calls — each its own `conn.begin()`, repo-scoped membership subtraction (never delete-by-path) and an exact-match repo purge respectively; neither decides WHAT is retired/desired, only applies a caller-supplied set (see `job.py`'s row). | | `symbols.py` | `extract_file` (#84): one tree-sitter parse, one full-tree walk, emitting both `symbols` and `reference_edges` candidates (`FileExtraction`) in the same pass — nested definitions captured, named nodes only for symbols, kinds from `SYMBOL_KINDS`/`EDGE_NODE_KINDS`, 1-based lines. Stack entries carry the innermost named enclosing definition so edges attribute to it in O(1) with no second walk. `extract_symbols` is a thin wrapper (`extract_file(pf).symbols`) kept for the existing call sites. Python-only edge helpers (`_python_call_target`, `_python_import_edges`) are the seam #85 generalizes to other languages. Parser cache is per-thread (`threading.local`) as insurance against a future GIL-releasing `parse()`. | +| `timing.py` | Per-phase wall-clock accounting for one branch (#103). `PhaseTimer` accumulates `phase -> seconds`; `install_timer`/`reset_timer`/`current_timer` carry one ambiently in a `ContextVar` (the same idiom as `job.py`'s `_repo_ctx`, for the same cross-module attribution problem); `record(phase, seconds)` is a **no-op when no timer is installed and never raises**, so `index_repo` stays callable directly. `_CLOCK` (default `time.monotonic`) is the single clock source every asserted duration reads — via `now()`, or via `PhaseTimer.clock` captured at construction — and is the seam tests patch; never read `time.monotonic()` directly for a number that appears on those lines. Stdlib only. | ## For AI Agents @@ -38,7 +39,7 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - **Import discipline**: `repo_config.py` stays import-light (no httpx/SQLAlchemy/SDK at module level); `normalize_repo` is imported from `repo_config`, never from `job` (import cycle); `app.embed`'s databricks-sdk dependency is only imported when semantic is on. ### Testing Requirements -- `make test` (`pytest -m "unit or observability"`, no external deps): `tests/unit/test_branches.py`, `test_chunk_store.py`, `test_chunking.py`, `test_edges.py`, `test_fetch.py`, `test_job.py`, `test_job_redaction.py`, `test_languages.py`, `test_parse.py`, `test_repo_config.py`, `test_resolve.py`, `test_symbols.py`, `test_store_chunk_writer.py`, `test_semantics_version_tripwire.py`. +- `make test` (`pytest -m "unit or observability"`, no external deps): `tests/unit/test_branches.py`, `test_chunk_store.py`, `test_chunking.py`, `test_edges.py`, `test_fetch.py`, `test_job.py`, `test_job_redaction.py`, `test_languages.py`, `test_parse.py`, `test_repo_config.py`, `test_resolve.py`, `test_symbols.py`, `test_store_chunk_writer.py`, `test_semantics_version_tripwire.py`, `test_timing.py`. - `make test-integration` (needs Postgres): `tests/integration/test_store.py`, `test_store_chunk_writer.py`, `test_content_sha_parity.py` (the hard gate on `hashing.py` vs the migration backfill), `test_reconcile.py` (the `reconcile_retired_branches`/`reconcile_removed_repos` storage primitives directly), and `test_job_reconcile.py` (#59: the real `run()` end to end — real primitives, real engine, only GitHub HTTP faked). - Schema tests must stay fast: changes to `repo_config.py` must not add heavy imports. diff --git a/indexer/job.py b/indexer/job.py index 83853be..3a0f1d8 100644 --- a/indexer/job.py +++ b/indexer/job.py @@ -79,6 +79,18 @@ loop as symbols. Flag-off: no chunking, no embedder, no import of ``app.embed``'s lazy ``databricks-sdk`` dependency. +Every INDEXED branch also emits one ``phase timing`` line accounting for its whole +wall clock -- resolve / download / extract / parse / embed / db / sweep, plus an +``other`` residual so no time is silently unattributed. The fields are fixed and +unconditional (a phase that did not run prints ``0.00s``) so the line stays +greppable whether or not semantic indexing is on. Skipped, failed, and conflicted +branches emit no such line. ``sweep`` is measured inside ``indexer.store`` and +reaches this module through the ambient timer in :mod:`indexer.timing` rather +than through the ``index_fn`` seam, which keeps that seam's signature (and every +fake of it) unchanged. The repo-level costs that sit outside every branch's +total -- the default branch's HEAD resolve and the branch listing -- are reported +as ``resolve=``/``list=`` on the per-repo ``finished`` line instead. + Logging is INFO only. The GitHub token is read via an injected client and is never logged, and this module never lowers root/SDK/httpx log levels (see the redaction test + source-level tripwire). @@ -93,7 +105,7 @@ import sys import tempfile import time -from collections.abc import Callable +from collections.abc import Callable, Iterable, Iterator from concurrent.futures import ThreadPoolExecutor, as_completed from contextvars import ContextVar from dataclasses import dataclass, field @@ -118,7 +130,7 @@ resolve_branch_head, resolve_ref, ) -from indexer.languages import Chunk, IndexCounts, ParsedFile +from indexer.languages import Chunk, FileExtraction, IndexCounts, ParsedFile from indexer.parse import iter_chunks, iter_source_files from indexer.repo_config import RepoConfig, effective_workers, load_config, normalize_repo from indexer.resolve import MAX_REPOS, RepoEntry, resolve_repos @@ -131,6 +143,7 @@ reconcile_retired_branches, ) from indexer.symbols import extract_file +from indexer.timing import PhaseTimer, install_timer, now, reset_timer logger = logging.getLogger("indexer.job") @@ -887,7 +900,15 @@ def _index_one_inner( """ name = normalize_repo(entry.name) org, repo = name.split("/", 1) + # Timed with two plain locals off the same clock the per-branch PhaseTimer + # reads -- no timer is installed here, because these costs are REPO-scoped + # and belong to no branch's total (they are reported on the `finished` line + # below instead). Without them the default branch's resolve would show as + # `resolve=0.00s` on its phase line, and a glob-configured monorepo's + # paginated branch listing would appear in no field on no line at all. + resolve_started = now() default_branch, default_head_sha = resolve_ref(http_client, org, repo) + resolve_elapsed = now() - resolve_started # An override matched to THIS repo by resolve_repos wins outright over the global # cap -- it is not a floor/ceiling blend, since a repo big enough to need one @@ -904,7 +925,9 @@ def _index_one_inner( # The common case -- no branches: configured -- needs no GitHub branches API # call at all: resolve_branches ignores all_branches entirely when globs is # empty (it always resolves to just [default_branch]). + list_started = now() all_branches = list_branches(http_client, org, repo) if entry.branch_globs else [] + list_elapsed = now() - list_started resolution = resolve_branches( default_branch, all_branches, sorted(entry.branch_globs), repo=name ) @@ -934,10 +957,48 @@ def _index_one_inner( # assertions, and a nonzero-by-construction timing field would break them. # This is the instrument the "throughput measured on the first production # run" promise depends on -- without it that promise is unfalsifiable. - logger.info("finished %s in %.2fs", name, time.monotonic() - started) + # `resolve=`/`list=` are the repo-scoped costs that appear in no branch's + # `phase timing` total. Printed unconditionally (`list=0.00s` whenever the + # repo has no `branches:` globs and the endpoint is never called) so the line + # never drifts in shape. The elapsed value keeps its own time.monotonic() + # reading -- it is not reconciled against the parenthesised numbers. + logger.info( + "finished %s in %.2fs (resolve=%.2fs list=%.2fs)", + name, + time.monotonic() - started, + resolve_elapsed, + list_elapsed, + ) return RepoOutcome(name=name, discovery_complete=resolution.complete, outcomes=outcomes) +def _timed_items( + items: Iterable[tuple[ParsedFile, FileExtraction]], timer: PhaseTimer +) -> Iterator[tuple[ParsedFile, FileExtraction]]: + """Charge each item's PRODUCTION to ``parse``, leaving its consumption to ``db``. + + ``items`` is consumed lazily inside ``index_repo``'s open transaction (the + bounded-memory invariant), so timing ``index_fn`` alone fuses parse and db + into one number, and materializing the generator to separate them would break + that invariant. This wrapper times only the ``next()`` calls -- the walk, the + read, the decode, and the tree-sitter extraction -- so the DML between them + stays attributable to ``db``. Never build a list here. + + The post-``StopIteration`` charge is not bookkeeping pedantry: on the + non-semantic path the final ``next()`` is where the directory walk finishes. + """ + it = iter(items) + while True: + t0 = timer.clock() + try: + item = next(it) + except StopIteration: + timer.add("parse", timer.clock() - t0) + return + timer.add("parse", timer.clock() - t0) + yield item + + def _index_one_branch( name: str, *, @@ -967,11 +1028,30 @@ def _index_one_branch( matched, else ``cfg.semantic_max_chunks_per_repo``. Taken as a parameter rather than read from ``cfg`` directly so every branch of a repo enforces the SAME resolved cap without recomputing (or risking drift on) the override lookup. + + Every phase of the indexed path is wall-clocked into ``timer`` and reported on + one ``phase timing`` line at the end (see the module docstring). ``timer`` is + also installed as the AMBIENT timer for the duration, which is how + ``indexer.store``'s sweep -- behind the ``index_fn`` seam, in another module -- + lands in this branch's numbers without changing that seam's signature. The + ``finally: reset_timer(token)`` is mandatory for the same reason + ``_index_one``'s ``_repo_ctx`` reset is: ``ThreadPoolExecutor`` reuses worker + threads without resetting their context, so a leaked timer would attribute + this branch's sweep to the next task on this thread. """ + # NOT the `started` parameter: that one is set once per REPO and passed + # unchanged to every branch, so reusing it would make branch 2+ of a + # glob-configured repo report a repo-cumulative total and silently inflate + # `other` by every preceding branch's wall clock. + timer = PhaseTimer() + branch_started = timer.clock() + timer_token = install_timer(timer) try: + t0 = timer.clock() head_sha = ( default_head_sha if is_default else resolve_branch_head(http_client, org, repo, branch) ) + timer.add("resolve", timer.clock() - t0) # The skip seam: after the immutable HEAD SHA is known, before anything # is downloaded. Both halves must match -- a stored NULL version never does. @@ -993,15 +1073,35 @@ def _index_one_branch( # being written to) and BEFORE the first byte is downloaded. Raising # here is caught below, costing this branch alone. assert_disk_headroom(tmp_path, repo=f"{name}@{branch}") + + t0 = timer.clock() tar_path = download_tarball(http_client, org, repo, head_sha, tmp_path) + timer.add("download", timer.clock() - t0) + + t0 = timer.clock() root = extract_tarball(tar_path, tmp_path / "extracted") + timer.add("extract", timer.clock() - t0) chunk_writer: ChunkWriter | None = None if cfg.semantic_enabled and embed_fn is not None: # Chunking/embedding needs the full file list up front -- unlike # the lazy items generator below, it cannot stream through # index_repo's open transaction. + # + # This walk is charged to `parse` explicitly: it is the same + # rglob + stat + read + decode work that _timed_items charges on + # the non-semantic path, and leaving it unwrapped would dump the + # entire file-walk cost of the PRODUCTION (semantic-on) path into + # `other`, which is exactly the number this instrumentation + # exists to route work by. + t0 = timer.clock() files = list(iter_source_files(root)) + timer.add("parse", timer.clock() - t0) + + # In a `finally`, unlike every other phase wrap: the degrade path + # below still burned this time (a downed embedder can burn a lot + # of it before it gives up) and must still be reported. + t0 = timer.clock() try: chunk_writer = _precompute_chunk_writer(files, embed_fn, max_chunks_per_repo) except Exception: @@ -1018,11 +1118,23 @@ def _index_one_branch( exc_info=True, ) chunk_writer = None + finally: + timer.add("embed", timer.clock() - t0) items = ((pf, extract_file(pf)) for pf in files) else: # Lazy generator: files stream through the open transaction (bounded memory). items = ((pf, extract_file(pf)) for pf in iter_source_files(root)) + # Parse time is INSIDE the db window (items are produced lazily as + # index_repo consumes them), so `db` subtracts only the parse accrued + # DURING that window -- never the phase total, which on the semantic + # path already holds the eager walk above. `sweep` is windowed the + # same way for the same reason (store.py records it from inside + # index_fn) -- kept symmetric with `parse` so a future sweep call + # site outside this window can't silently mis-window `db`. + parse_before = timer.total("parse") + sweep_before = timer.total("sweep") + t0 = timer.clock() with engine.connect() as conn: counts = index_fn( conn, @@ -1030,9 +1142,62 @@ def _index_one_branch( branch=branch, is_default=is_default, head_sha=head_sha, - items=items, + items=_timed_items(items, timer), chunk_writer=chunk_writer, ) + db_wall = timer.clock() - t0 + + # Emitted HERE -- after the TemporaryDirectory teardown (an rm -rf of a + # possibly multi-GB extracted tree, which `other` must include) and from + # inside the worker, where _index_one's _repo_ctx still resolves the + # [%(repo)s] field. The drain loop on the main thread would render `[-]`. + # + # `index_fn` above already committed this branch's transaction -- `counts` + # is proof of that. This block is measurement ONLY, so it gets its own + # try/except: per timing.py's own principle ("instrumentation must never + # be able to fail the work it measures"), a bug here must degrade to a + # missing log line, never to reclassifying an already-committed branch as + # `failed` (which would also flip the run's exit code and gate off the + # post-fan-out reconciliation checkpoint, which requires zero failures). + try: + db = max( + 0.0, + db_wall + - (timer.total("parse") - parse_before) + - (timer.total("sweep") - sweep_before), + ) + total = timer.clock() - branch_started + other = max( + 0.0, + total + - timer.total("resolve") + - timer.total("download") + - timer.total("extract") + - timer.total("parse") + - timer.total("embed") + - db + - timer.total("sweep"), + ) + # One format string, no branches: a phase that did not run prints 0.00s + # rather than vanishing, so the line stays greppable and field-stable + # whether or not semantic indexing is on. + logger.info( + "phase timing %s@%s: total=%.2fs resolve=%.2fs download=%.2fs extract=%.2fs " + "parse=%.2fs embed=%.2fs db=%.2fs sweep=%.2fs other=%.2fs", + name, + branch, + total, + timer.total("resolve"), + timer.total("download"), + timer.total("extract"), + timer.total("parse"), + timer.total("embed"), + db, + timer.total("sweep"), + other, + ) + except Exception: + logger.warning("phase timing unavailable for %s@%s", name, branch, exc_info=True) return BranchOutcome(branch=branch, status="indexed", counts=counts) except StaleIndexError as exc: # The repo_branches row for THIS branch changed under this worker, so @@ -1065,6 +1230,12 @@ def _index_one_branch( except Exception: logger.exception("failed to index %s@%s", name, branch) return BranchOutcome(branch=branch, status="failed") + finally: + # Mandatory, exactly like _index_one's _repo_ctx reset: ThreadPoolExecutor + # reuses worker threads without resetting their context, so a leaked timer + # would attribute this branch's sweep to the NEXT branch (or repo) that + # lands on this thread -- silently wrong numbers, which is worse than none. + reset_timer(timer_token) def _positive_int(raw: str) -> int: diff --git a/indexer/store.py b/indexer/store.py index f85f0df..84e110a 100644 --- a/indexer/store.py +++ b/indexer/store.py @@ -29,6 +29,7 @@ from app.db.models import INDEX_SEMANTICS_VERSION, File, ReferenceEdge, Repo, RepoBranch, Symbol from indexer.hashing import content_sha from indexer.languages import FileExtraction, IndexCounts, ParsedFile +from indexer.timing import now, record logger = logging.getLogger("indexer.store") @@ -242,6 +243,14 @@ def index_repo( if chunk_writer is not None: chunk_writer(conn, repo_id, file_id, pf) + # Timed into indexer.job's ambient per-branch PhaseTimer, if one is + # installed -- a no-op otherwise, so a direct index_repo call (tests, + # scripts) is unaffected. Deliberately NOT a return value: IndexCounts is + # a frozen dataclass compared by value in existing assertions, and NOT a + # new index_repo parameter: that signature is an injected seam whose + # fakes would all have to grow one. No log record is emitted here; the + # number surfaces on job.py's single `phase timing` line. + sweep_started = now() swept = _sweep_membership( conn, name=name, @@ -250,6 +259,7 @@ def index_repo( seen_paths=seen_paths, seen_shas=seen_shas, ) + record("sweep", now() - sweep_started) _stamp_repo_branch( conn, diff --git a/indexer/timing.py b/indexer/timing.py new file mode 100644 index 0000000..4a7f477 --- /dev/null +++ b/indexer/timing.py @@ -0,0 +1,117 @@ +"""Ambient per-phase wall-clock accounting for one branch's indexing pipeline. + +``indexer.job`` measures most of a branch's phases (resolve / download / extract / +parse / embed / db) inline, but ``sweep`` runs deep inside +:func:`indexer.store.index_repo`, behind the injected ``index_fn`` seam and a +frozen ``IndexCounts`` return type. Threading a timer through that seam would +force every existing ``index_fn`` fake to grow a parameter, turning a log-only +change into a storage-contract change -- so the timer is carried *ambiently* in a +:class:`contextvars.ContextVar`, exactly mirroring ``indexer.job``'s ``_repo_ctx`` ++ ``RepoLogFilter`` idiom for the same cross-module attribution problem. + +Two properties are load-bearing: + +* **Default-unset is a no-op.** :func:`record` with no installed timer does + nothing and never raises, so ``index_repo`` stays callable directly (as + ``tests/integration/test_store.py`` does) with no timer in sight. +* **Install/reset discipline.** ``ThreadPoolExecutor`` reuses worker threads and + does NOT reset their context between tasks, so :func:`install_timer` must + always be paired with :func:`reset_timer` in a ``finally`` -- a leaked timer + would silently attribute one branch's sweep to the next branch's line. + +The ``ContextVar`` does not cross a thread or process boundary: it isolates +concurrent worker threads from each other (which is the property this module +needs today), but a call to :func:`record` from a different thread or a +``ProcessPoolExecutor`` worker than the one that called :func:`install_timer` +reaches a *different* ambient timer (or none), not the installing thread's. +Relevant if a future phase's work is moved off the indexing worker thread -- +e.g. issue #108's process-pool extraction, or issue #107's concurrent +embedding. + +``_CLOCK`` is the module's single clock source. Every interval that appears on +an asserted log line reads it (via :func:`now`, or via +:attr:`PhaseTimer.clock` which captures it at construction), never a bare +``time.monotonic()`` -- mixing a fake clock with the real one makes the residual +``other=`` field meaningless. Tests patch ``indexer.timing._CLOCK`` before +driving ``run()``; that is the only injection point that reaches a +worker-constructed timer. + +Stdlib only (``time`` + ``contextvars``): this module is imported by the +indexing hot path and adds no dependency. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from contextvars import ContextVar, Token + +# The single clock source for every duration on the phase-timing and repo +# `finished` log lines. Monotonic (never NTP-adjusted) and patched wholesale by +# tests -- do not read time.monotonic() directly anywhere those numbers are +# asserted against each other. +_CLOCK: Callable[[], float] = time.monotonic + + +def now() -> float: + """Read the current clock, resolving ``_CLOCK`` at call time (test seam).""" + return _CLOCK() + + +class PhaseTimer: + """A per-branch accumulator of ``phase -> seconds``. + + Not thread-safe by design: one timer belongs to exactly one branch, which is + indexed by exactly one worker thread. Cross-thread isolation comes from the + ``ContextVar`` below, not from locking. + """ + + def __init__(self) -> None: + # Captured at construction (not read per call) so one branch's + # arithmetic can never straddle a clock swap mid-flight. There is no + # constructor override for this: `record()` reaches this timer only + # through the ambient ContextVar, always via the module-level + # `_CLOCK`, so a per-instance clock would silently desync from the + # sweep timing recorded through `now()` in indexer/store.py. + self.clock: Callable[[], float] = _CLOCK + self._totals: dict[str, float] = {} + + def add(self, phase: str, seconds: float) -> None: + """Accumulate ``seconds`` into ``phase`` (phases are accrued, not set).""" + self._totals[phase] = self._totals.get(phase, 0.0) + seconds + + def total(self, phase: str) -> float: + """Seconds accrued to ``phase``; ``0.0`` for a phase that never ran.""" + return self._totals.get(phase, 0.0) + + +# Ambient per-thread timer. Default None = "nobody is measuring", which is the +# state every direct index_repo caller (and every non-worker thread) sees. +_timer_ctx: ContextVar[PhaseTimer | None] = ContextVar("phase_timer", default=None) + + +def install_timer(timer: PhaseTimer) -> Token[PhaseTimer | None]: + """Make ``timer`` the ambient one for this context; reset the token in a ``finally``.""" + return _timer_ctx.set(timer) + + +def reset_timer(token: Token[PhaseTimer | None]) -> None: + """Undo :func:`install_timer`. Mandatory -- see the module docstring.""" + _timer_ctx.reset(token) + + +def current_timer() -> PhaseTimer | None: + """The ambient timer, or ``None`` when nothing is measuring.""" + return _timer_ctx.get() + + +def record(phase: str, seconds: float) -> None: + """Accrue ``seconds`` to ``phase`` on the ambient timer, if there is one. + + A no-op when no timer is installed. Instrumentation must never be able to + fail the work it measures. + """ + timer = _timer_ctx.get() + if timer is None: + return + timer.add(phase, seconds) diff --git a/tests/integration/test_store.py b/tests/integration/test_store.py index 194fef5..22cc75e 100644 --- a/tests/integration/test_store.py +++ b/tests/integration/test_store.py @@ -32,6 +32,7 @@ ParsedFile, ) from indexer.store import StaleIndexError, _stamp_repo_branch, index_repo +from indexer.timing import PhaseTimer, install_timer, reset_timer SCHEMA = "test_store" @@ -178,6 +179,38 @@ def test_mark_and_sweep_removes_deleted_file(conn: Connection) -> None: assert _count(conn, "files", "commit = 'sha_second'") == 1 +@pytest.mark.integration +def test_index_repo_records_the_sweep_phase(conn: Connection) -> None: + """The sweep duration reaches ``indexer.job``'s per-branch line ambiently. + + ``index_repo`` returns a frozen ``IndexCounts`` and takes no timer parameter -- + the sweep's cost travels out through the ``indexer.timing`` ContextVar instead, + which is what keeps the ``index_fn`` seam (and every fake of it) unchanged. + ``tests/unit/test_job.py`` pins the arithmetic against a fake ``index_fn``; + this pins the real cross-module wiring against the real sweep and real SQL. + """ + _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(MAIN, UTIL)) + conn.rollback() + + timer = PhaseTimer() + token = install_timer(timer) + try: + # Re-run without util.py at a new SHA: the same scenario as + # test_mark_and_sweep_removes_deleted_file, so the counts are unchanged. + counts = _index_default( + conn, name="acme/widgets", head_sha="sha_second", items=_items(MAIN) + ) + finally: + reset_timer(token) + + assert counts == IndexCounts(files=1, symbols=1, swept=1, edges=0) + assert timer.total("sweep") > 0.0 + # index_repo measures the sweep and nothing else -- every other phase is + # job.py's to record. + assert timer.total("db") == 0.0 + assert timer.total("parse") == 0.0 + + @pytest.mark.integration def test_sweep_is_repo_scoped(conn: Connection) -> None: # Two repos indexed at their own SHAs; re-indexing repo A with a new SHA must diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 85c1b30..d4ca8dd 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -15,6 +15,7 @@ from __future__ import annotations import base64 +import dataclasses import inspect import io import logging @@ -2266,3 +2267,732 @@ def test_reconcile_unit_level_partial_progress_on_mid_sequence_failure() -> None assert progress.committed_any is True assert progress.purge_blocked is False assert progress.purged_repos == [] + + +# --- per-phase timing instrumentation (#103) -------------------------------- +# Every INDEXED branch emits one `phase timing` record accounting for its whole +# wall clock. The numbers are asserted EXACTLY, never with sleeps: the tests +# below patch `indexer.timing._CLOCK` (the single seam every asserted interval +# reads, captured by PhaseTimer.__init__ inside the worker) with a counter that +# only moves when a fake explicitly moves it. A sleep-based version of these +# assertions would be flaky under CI load and could not pin a residual at all. +# +# `_CLOCK` is a module global, so a fake counter is shared process-wide while a +# run fans out. Any fake-clock test driving more than one repo therefore pins +# `index_concurrency=1`, which gives one worker thread and one deterministic +# advance sequence. + +# Anchored and exhaustive on purpose: this single constant pins the field set, +# the field ORDER, the `%.2f` shape, the mandatory single-space separator, and +# the absence of anything else on the line (no config values, no counts, no +# headers -- the redaction posture). `\d+\.\d\d` also rejects a negative value +# outright, which is what makes the `max(0.0, ...)` clamps in job.py testable. +_TIMING_RE = re.compile( + r"^phase timing (?P[^ @]+)@(?P\S+): " + r"total=(\d+\.\d\d)s resolve=(\d+\.\d\d)s download=(\d+\.\d\d)s " + r"extract=(\d+\.\d\d)s parse=(\d+\.\d\d)s embed=(\d+\.\d\d)s " + r"db=(\d+\.\d\d)s sweep=(\d+\.\d\d)s other=(\d+\.\d\d)s$" +) + + +class _FakeClock: + """A monotonic clock that advances only when a test advances it.""" + + def __init__(self) -> None: + self.t = 0.0 + + def __call__(self) -> float: + return self.t + + def advance(self, seconds: float) -> None: + self.t += seconds + + +def _install_fake_clock(monkeypatch: pytest.MonkeyPatch) -> _FakeClock: + """Patch the one clock seam every asserted duration reads.""" + clock = _FakeClock() + monkeypatch.setattr("indexer.timing._CLOCK", clock) + return clock + + +def _timing_messages(caplog: pytest.LogCaptureFixture) -> list[str]: + return [m for r in caplog.records if (m := r.getMessage()).startswith("phase timing ")] + + +def _phases(message: str) -> dict[str, float]: + """Parse one `phase timing` line into ``{field: seconds}``, pinning its shape first.""" + assert _TIMING_RE.match(message), f"malformed timing line: {message!r}" + body = message.split(": ", 1)[1] + return {k: float(v.removesuffix("s")) for k, v in (f.split("=") for f in body.split(" "))} + + +def _only_phases(caplog: pytest.LogCaptureFixture) -> dict[str, float]: + """The single expected timing line's fields.""" + messages = _timing_messages(caplog) + assert len(messages) == 1, f"expected exactly one timing line, got {messages}" + return _phases(messages[0]) + + +@pytest.mark.unit +def test_indexed_branch_logs_every_phase_field(caplog: pytest.LogCaptureFixture) -> None: + """T1: one record per indexed branch, with all nine fields in a fixed order.""" + idx = _RecordingIndex() + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), idx) + assert code == 0 + + messages = _timing_messages(caplog) + assert len(messages) == 1 + match = _TIMING_RE.match(messages[0]) + assert match is not None, messages[0] + assert match.group("repo") == "acme/widgets" + assert match.group("branch") == "main" + assert list(_phases(messages[0])) == [ + "total", + "resolve", + "download", + "extract", + "parse", + "embed", + "db", + "sweep", + "other", + ] + + +@pytest.mark.unit +def test_semantic_off_still_logs_embed_zero(caplog: pytest.LogCaptureFixture) -> None: + """T2 (AC2): the format string has no branches -- semantic off zeroes embed, it does + not drop the field. A conditional line would break every grep an operator writes.""" + cfg_off = Settings(semantic_enabled=False) + with caplog.at_level(logging.INFO, logger="indexer.job"): + assert _run(_config(repos=["acme/widgets"]), _RecordingIndex(), cfg=cfg_off) == 0 + off = _only_phases(caplog) + assert off["embed"] == 0.0 + + caplog.clear() + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + _config(repos=["acme/widgets"]), + _RecordingIndex(), + cfg=cfg, + embed_fn=lambda texts: [[0.0] for _ in texts], + ) + assert code == 0 + on = _only_phases(caplog) + + # Byte-identical field set and order, semantic on or off. + assert list(on) == list(off) + + +@pytest.mark.unit +def test_semantic_on_attributes_embed_time( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T3: chunking + embedding is its own phase, and it is NOT double-counted into db.""" + clock = _install_fake_clock(monkeypatch) + + def _embed(texts: list[str]) -> list[list[float]]: + clock.advance(4.0) + return [[0.0] for _ in texts] + + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), _RecordingIndex(), cfg=cfg, embed_fn=_embed) + assert code == 0 + + phases = _only_phases(caplog) + assert phases["embed"] == 4.0 + assert phases["db"] == 0.0 # the embed happens BEFORE engine.connect(), not inside it + assert phases["total"] == 4.0 + assert phases["other"] == 0.0 + + +@pytest.mark.unit +def test_phase_timing_line_carries_the_repo_context(caplog: pytest.LogCaptureFixture) -> None: + """T4 (AC1): emitted from inside the worker, so `[%(repo)s]` resolves. + + The drain loop's `indexed ...` line runs on the main thread, where the context + var is at its default and the record would render `[-]` -- which is exactly + why the timing line is NOT appended there. + """ + import indexer.job as job + + log_filter = job.RepoLogFilter() + caplog.handler.addFilter(log_filter) + try: + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), _RecordingIndex()) + finally: + caplog.handler.removeFilter(log_filter) + assert code == 0 + + timing_records = [r for r in caplog.records if r.getMessage().startswith("phase timing ")] + assert len(timing_records) == 1 + assert timing_records[0].repo == "acme/widgets" + + +@pytest.mark.unit +def test_parse_time_is_excluded_from_db_time( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T5: items are produced lazily INSIDE index_repo's transaction, so a naive + "time index_fn" would report parse and db fused. _timed_items charges only the + production of each item to parse, leaving the DML between them to db.""" + import indexer.job as job + + clock = _install_fake_clock(monkeypatch) + real_extract_file = job.extract_file + + def _slow_extract(pf: ParsedFile) -> Any: + clock.advance(2.0) + return real_extract_file(pf) + + monkeypatch.setattr(job, "extract_file", _slow_extract) + + def _index(conn: Any, *, items: Any, **_: Any) -> IndexCounts: + consumed = list(items) + clock.advance(9.0) + return IndexCounts(files=len(consumed), symbols=0, swept=0, edges=0) + + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), _index) + assert code == 0 + + phases = _only_phases(caplog) + assert phases["parse"] == 4.0 # main.py + README.md, 2.0s each + assert phases["db"] == 9.0 # NOT 13.0 + assert phases["total"] == 13.0 + assert phases["other"] == 0.0 + + +@pytest.mark.unit +def test_skipped_branch_logs_no_timing_line(caplog: pytest.LogCaptureFixture) -> None: + """T6: AC1 says "every indexed (not skipped) branch". A skipped branch has no + phases to report and must not pollute the dominant-phase grep with zeroes.""" + engine = _FakeEngine( + stamps={("acme/widgets", "main"): ("sha_widgets", INDEX_SEMANTICS_VERSION)} + ) + idx = _RecordingIndex() + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), idx, engine=engine) + assert code == 0 + assert idx.calls == [] + assert "skipped acme/widgets@main" in caplog.text + assert _timing_messages(caplog) == [] + + +@pytest.mark.unit +def test_failed_branch_logs_no_timing_line(caplog: pytest.LogCaptureFixture) -> None: + """T7: partial timings for a branch that never finished would be misleading.""" + + def _boom(conn: Any, *, items: Any, **_: Any) -> IndexCounts: + raise RuntimeError("boom") + + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), _boom) + assert code == 1 + assert "failed to index acme/widgets@main" in caplog.text + assert _timing_messages(caplog) == [] + + +@pytest.mark.unit +def test_conflicted_branch_logs_no_timing_line(caplog: pytest.LogCaptureFixture) -> None: + """T20: the StaleIndexError path is classified separately from T7's generic + failure, and it too rolled the branch back -- nothing was indexed, so nothing + is reported.""" + + def _stale(conn: Any, *, name: str, branch: str, items: Any, **_: Any) -> IndexCounts: + raise StaleIndexError(f"repo_branches row for {name}@{branch} changed") + + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), _stale) + assert code == 0 # a conflict self-heals and does not fail the run + assert "index conflict for acme/widgets@main" in caplog.text + assert _timing_messages(caplog) == [] + + +@pytest.mark.unit +def test_repo_line_reports_resolve_and_branch_listing( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T8: the two REPO-scoped costs belong to no branch's total. + + Without `resolve=` the default branch's HEAD resolve would show as + `resolve=0.00s` on its phase line and look like a bug; without `list=` a + glob-configured monorepo's paginated branch listing would appear in no field + on no line at all. + + Deliberately asserts NO relation between the `in ...s` elapsed and the + parenthesised values: the former keeps its own real `time.monotonic()` + reading, so under a fake clock it reads ~0.00s while the others do not. + """ + import indexer.job as job + + clock = _install_fake_clock(monkeypatch) + real_resolve_ref = job.resolve_ref + real_list_branches = job.list_branches + + def _resolve_ref(*args: Any, **kwargs: Any) -> tuple[str, str]: + clock.advance(5.0) + return real_resolve_ref(*args, **kwargs) + + def _list_branches(*args: Any, **kwargs: Any) -> list[str]: + clock.advance(3.0) + return real_list_branches(*args, **kwargs) + + monkeypatch.setattr(job, "resolve_ref", _resolve_ref) + monkeypatch.setattr(job, "list_branches", _list_branches) + + # No `branches:` globs -> the branches endpoint is never called at all. + with caplog.at_level(logging.INFO, logger="indexer.job"): + assert _run(_config(repos=["acme/widgets"]), _RecordingIndex()) == 0 + finished = re.search( + r"finished acme/widgets in ([0-9.]+)s \(resolve=([0-9.]+)s list=([0-9.]+)s\)", caplog.text + ) + assert finished is not None, caplog.text + assert float(finished.group(2)) == 5.0 + assert float(finished.group(3)) == 0.0 + # The pre-existing unanchored assertion on this line still matches. + assert re.search(r"finished acme/widgets in ([0-9.]+)s", caplog.text) is not None + + caplog.clear() + github = _GitHub(branches={"acme/widgets": ["main"]}) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + _config(repos=["acme/widgets"], branches=["*"], index_concurrency=1), + _RecordingIndex(), + github=github, + ) + assert code == 0 + globbed = re.search( + r"finished acme/widgets in ([0-9.]+)s \(resolve=([0-9.]+)s list=([0-9.]+)s\)", caplog.text + ) + assert globbed is not None, caplog.text + assert float(globbed.group(2)) == 5.0 + assert float(globbed.group(3)) == 3.0 + # ...and none of it leaked into the branch's own total. + assert _only_phases(caplog)["total"] == 0.0 + + +@pytest.mark.unit +def test_timings_do_not_leak_between_branches_of_one_repo( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T9: branch 2's numbers are its own, not cumulative. + + The trap this guards is reusing `_index_one_branch`'s `started` parameter for + `total=`: that value is set once per REPO and passed unchanged to every + branch, so branch 2 would report branch 1's wall clock too and silently + inflate `other`. Also guards the ContextVar reset between branches. + """ + import indexer.job as job + + clock = _install_fake_clock(monkeypatch) + real_download = job.download_tarball + + def _download(*args: Any, **kwargs: Any) -> Any: + clock.advance(5.0) + return real_download(*args, **kwargs) + + monkeypatch.setattr(job, "download_tarball", _download) + + github = _GitHub(branches={"acme/widgets": ["main", "feature"]}) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + _config(repos=["acme/widgets"], branches=["feature"], index_concurrency=1), + _RecordingIndex(), + github=github, + ) + assert code == 0 + + messages = _timing_messages(caplog) + assert len(messages) == 2 + branches = {_TIMING_RE.match(m).group("branch") for m in messages} # type: ignore[union-attr] + assert branches == {"main", "feature"} + for message in messages: + phases = _phases(message) + assert phases["download"] == 5.0 + assert phases["total"] == 5.0 # not 10.0 on the second branch + assert phases["other"] == 0.0 + + +@pytest.mark.unit +def test_timings_do_not_leak_between_repos_on_a_reused_worker_thread( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T10: ThreadPoolExecutor reuses worker threads without resetting their + context. With one worker and two repos, repo 2 runs on the same thread that + just finished repo 1 -- a leaked timer would show up here as doubled numbers.""" + import indexer.job as job + + clock = _install_fake_clock(monkeypatch) + real_download = job.download_tarball + + def _download(*args: Any, **kwargs: Any) -> Any: + clock.advance(3.0) + return real_download(*args, **kwargs) + + monkeypatch.setattr(job, "download_tarball", _download) + + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + _config(repos=["acme/widgets", "acme/gadgets"], index_concurrency=1), + _RecordingIndex(), + ) + assert code == 0 + + messages = _timing_messages(caplog) + assert len(messages) == 2 + assert {_TIMING_RE.match(m).group("repo") for m in messages} == { # type: ignore[union-attr] + "acme/widgets", + "acme/gadgets", + } + for message in messages: + phases = _phases(message) + assert phases["download"] == 3.0 + assert phases["total"] == 3.0 + assert phases["other"] == 0.0 + + +@pytest.mark.unit +def test_timer_reset_fires_on_every_branch_exit_path(monkeypatch: pytest.MonkeyPatch) -> None: + """Directly guards the `finally: reset_timer(token)` invariant itself. + + T9/T10 assert branch totals aren't cumulative, but `_index_one_branch` + unconditionally does `install_timer(PhaseTimer())` at entry -- which + shadows a prior leaked timer regardless of whether `reset_timer` actually + ran. So T9/T10 pass identically whether or not the `finally` fires, and + cannot detect its removal. This test spies on `indexer.timing.reset_timer` + directly and asserts it is called exactly once per branch attempt, + including the failed and conflicted paths -- the actual invariant the + module docstring calls "mandatory". + """ + import indexer.job as job + import indexer.timing as timing + + calls = 0 + real_reset = timing.reset_timer + + def _spy_reset(token: Any) -> None: + nonlocal calls + calls += 1 + real_reset(token) + + monkeypatch.setattr(job, "reset_timer", _spy_reset) + + def _ok(conn: Any, *, items: Any, **_: Any) -> IndexCounts: + return IndexCounts(files=len(list(items)), symbols=0, swept=0, edges=0) + + def _boom(conn: Any, *, items: Any, **_: Any) -> IndexCounts: + raise RuntimeError("boom") + + def _stale(conn: Any, *, name: str, branch: str, items: Any, **_: Any) -> IndexCounts: + raise StaleIndexError(f"repo_branches row for {name}@{branch} changed") + + _run(_config(repos=["acme/widgets"]), _ok) + assert calls == 1 + + _run(_config(repos=["acme/widgets"]), _boom) + assert calls == 2 + + _run(_config(repos=["acme/widgets"]), _stale) + assert calls == 3 + + +@pytest.mark.unit +def test_phase_timing_failure_does_not_fail_an_already_committed_branch( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """A bug in the phase-timing arithmetic/log call must degrade to a missing + log line, never to reclassifying an already-committed branch as `failed`. + + By the point this block runs, `index_fn` has already returned `counts` -- + this branch's transaction is committed. Flipping it to `failed` here would + report a real committed index as a failure, flip the run's exit code, and + gate off `_decide_reconciliation` (which requires zero failures), all for a + bug in pure measurement code. Forces the failure at the `logger.info` call + itself (the last statement in the guarded block) so a passing test proves + the whole block is covered, not just the arithmetic above it. + """ + import indexer.job as job + + real_info = job.logger.info + + def _boom_on_phase_timing(msg: object, *args: Any, **kwargs: Any) -> None: + if isinstance(msg, str) and msg.startswith("phase timing "): + raise RuntimeError("boom") + real_info(msg, *args, **kwargs) + + monkeypatch.setattr(job.logger, "info", _boom_on_phase_timing) + + idx = _RecordingIndex() + with caplog.at_level(logging.WARNING, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), idx) + + assert code == 0 + assert idx.calls == ["acme/widgets"] + # No well-formed `phase timing ...: total=...` line was emitted (the raise + # happened inside the logger.info call itself) -- but the module's own + # "phase timing unavailable" fallback below shares the "phase timing " + # prefix, so check the anchored data-line shape rather than the prefix. + assert not any(_TIMING_RE.match(r.getMessage()) for r in caplog.records) + assert "failed to index" not in caplog.text + assert "phase timing unavailable for acme/widgets@main" in caplog.text + + +@pytest.mark.unit +def test_index_counts_fields_are_unchanged() -> None: + """T11 (AC3): a tripwire that timing never leaks into the frozen return type. + + `IndexCounts` is compared by value in existing assertions across the unit and + integration suites; a timing field would be nonzero by construction and break + every one of them. The sweep duration reaches job.py ambiently for exactly + this reason. + """ + assert [f.name for f in dataclasses.fields(IndexCounts)] == [ + "files", + "symbols", + "swept", + "edges", + ] + + +@pytest.mark.unit +def test_timing_line_contains_no_values_beyond_durations( + caplog: pytest.LogCaptureFixture, +) -> None: + """T12 (AC3 / redaction): the anchored regex matches the WHOLE message, so + nothing but the repo, the branch, and the nine durations can ever appear -- + no config values, no counts, no headers -- and no field can go negative.""" + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + config = _config(repos=["acme/widgets"], semantic_max_chunks_per_repo={"acme/widgets": 7}) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + config, + _RecordingIndex(), + cfg=cfg, + embed_fn=lambda texts: [[0.0] for _ in texts], + ) + assert code == 0 + + messages = _timing_messages(caplog) + assert len(messages) == 1 + assert _TIMING_RE.fullmatch(messages[0]) is not None, messages[0] + + +@pytest.mark.unit +def test_sweep_time_is_subtracted_from_db_time( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T16: the cross-module sweep hook, pinned without Postgres. + + A fake `index_fn` runs on the worker thread inside the installed timer's + context, so its `record("sweep", ...)` reaches exactly the timer the real + `store.py` hook would. Without this test the subtraction term is zero in + every unit test and the arithmetic is unproven under `make test`. + """ + import indexer.timing as timing + + clock = _install_fake_clock(monkeypatch) + + def _index(conn: Any, *, items: Any, **_: Any) -> IndexCounts: + consumed = list(items) + clock.advance(2.0) # the sweep really burned this time... + timing.record("sweep", 2.0) # ...and store.py reports it ambiently + clock.advance(6.0) # the rest of the transaction + return IndexCounts(files=len(consumed), symbols=0, swept=1, edges=0) + + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), _index) + assert code == 0 + + phases = _only_phases(caplog) + assert phases["sweep"] == 2.0 + assert phases["db"] == 6.0 # NOT 8.0 + assert phases["total"] == 8.0 + assert phases["other"] == 0.0 + + +@pytest.mark.unit +def test_semantic_eager_walk_is_charged_to_parse( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T18: on the semantic (production) path the file list is materialized up front, + OUTSIDE the db window and outside _timed_items. Leaving that walk unwrapped would + dump the whole rglob + stat + read + decode cost into `other` -- the exact number + that routes work between the epic's parse/IO/db issues.""" + import indexer.job as job + + clock = _install_fake_clock(monkeypatch) + real_iter_source_files = job.iter_source_files + + def _slow_walk(root: Any) -> Any: + # iter_source_files is a generator function: paying the cost at call + # time (before the first `next()`) would let this fake pass even if + # the real wrap only timed the call and not the iteration -- the + # exact mis-scoped-wrap bug this test exists to catch. Pay per + # yielded file instead, matching where the real cost (rglob walk, + # stat, read, decode) is actually incurred. + for pf in real_iter_source_files(root): + clock.advance(3.0) # 2 files in _DEFAULT_FILES -> 6.0s total + yield pf + + monkeypatch.setattr(job, "iter_source_files", _slow_walk) + + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + _config(repos=["acme/widgets"]), + _RecordingIndex(), + cfg=cfg, + embed_fn=lambda texts: [[0.0] for _ in texts], + ) + assert code == 0 + + phases = _only_phases(caplog) + assert phases["parse"] == 6.0 + assert phases["other"] == 0.0 # not swallowed by the residual + assert phases["db"] == 0.0 # and not misattributed to the transaction + assert phases["total"] == 6.0 + + +@pytest.mark.unit +def test_embed_degrade_path_still_logs_a_timing_line( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T19: a downed embedder degrades to a core index -- and it burned real time + getting there, so `embed=` must still report it. That is why the embed wrap is + the one phase recorded in a `finally`.""" + clock = _install_fake_clock(monkeypatch) + + def _down(_texts: list[str]) -> list[list[float]]: + clock.advance(7.0) + raise RuntimeError("serving endpoint unavailable") + + idx = _RecordingIndex() + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), idx, cfg=cfg, embed_fn=_down) + assert code == 0 + assert idx.calls == ["acme/widgets"] # the core index still ran + assert idx.chunk_writer is None # ...without chunks + + phases = _only_phases(caplog) + assert phases["embed"] == 7.0 + assert phases["total"] == 7.0 + assert phases["other"] == 0.0 + + +@pytest.mark.unit +def test_other_is_the_exact_unattributed_residual( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T17: pins the RESIDUAL, not the sum identity. + + `other` is *defined* as `total - sum(phases)`, so asserting + `total == sum(phases) + other` is a tautology that holds for any values and can + never detect a mis-measured phase. This test instead drives a distinct known + amount through every one of the seven phases and requires `other == 0.00s`, + then re-runs with one deliberately UNINSTRUMENTED advance and requires `other` + to equal exactly that amount. + """ + import indexer.job as job + + clock = _install_fake_clock(monkeypatch) + real_resolve_branch_head = job.resolve_branch_head + real_download_tarball = job.download_tarball + real_extract_tarball = job.extract_tarball + real_iter_source_files = job.iter_source_files + real_assert_disk_headroom = job.assert_disk_headroom + + def _resolve_branch_head(*args: Any, **kwargs: Any) -> str: + clock.advance(1.0) + return real_resolve_branch_head(*args, **kwargs) + + def _download_tarball(*args: Any, **kwargs: Any) -> Any: + clock.advance(2.0) + return real_download_tarball(*args, **kwargs) + + def _extract_tarball(*args: Any, **kwargs: Any) -> Any: + clock.advance(4.0) + return real_extract_tarball(*args, **kwargs) + + def _iter_source_files(root: Any) -> Any: + # Pay per yielded file, not at call time -- iter_source_files is a + # generator function, so a call-time advance would pass even for a + # wrap that only times the call and not the iteration. + for pf in real_iter_source_files(root): + clock.advance(4.0) # 2 files in _DEFAULT_FILES -> 8.0s total + yield pf + + def _embed(texts: list[str]) -> list[list[float]]: + clock.advance(16.0) + return [[0.0] for _ in texts] + + def _index(conn: Any, *, items: Any, **_: Any) -> IndexCounts: + import indexer.timing as timing + + consumed = list(items) + clock.advance(32.0) # the transaction's own DML + clock.advance(64.0) # the sweep, which store.py reports ambiently + timing.record("sweep", 64.0) + return IndexCounts(files=len(consumed), symbols=0, swept=0, edges=0) + + monkeypatch.setattr(job, "resolve_branch_head", _resolve_branch_head) + monkeypatch.setattr(job, "download_tarball", _download_tarball) + monkeypatch.setattr(job, "extract_tarball", _extract_tarball) + monkeypatch.setattr(job, "iter_source_files", _iter_source_files) + + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + # "main" is stamped current so it is SKIPPED, leaving exactly one indexed + # branch -- and a non-default one, which is the only way `resolve=` is + # non-zero (the default branch's HEAD came from the repo-level resolve_ref). + engine = _FakeEngine( + stamps={("acme/widgets", "main"): ("sha_widgets", INDEX_SEMANTICS_VERSION)} + ) + github = _GitHub(branches={"acme/widgets": ["main", "feature"]}) + config = _config(repos=["acme/widgets"], branches=["feature"], index_concurrency=1) + + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(config, _index, cfg=cfg, embed_fn=_embed, github=github, engine=engine) + assert code == 0 + + phases = _only_phases(caplog) + assert phases == { + "total": 127.0, + "resolve": 1.0, + "download": 2.0, + "extract": 4.0, + "parse": 8.0, + "embed": 16.0, + "db": 32.0, + "sweep": 64.0, + "other": 0.0, + } + + # Second leg: advance the clock at a point NOTHING instruments (the pre-flight + # disk check). Every phase must be unchanged and `other` must be exactly it. + def _assert_disk_headroom(*args: Any, **kwargs: Any) -> None: + clock.advance(9.0) + real_assert_disk_headroom(*args, **kwargs) + + monkeypatch.setattr(job, "assert_disk_headroom", _assert_disk_headroom) + caplog.clear() + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(config, _index, cfg=cfg, embed_fn=_embed, github=github, engine=engine) + assert code == 0 + + phases = _only_phases(caplog) + assert phases == { + "total": 136.0, + "resolve": 1.0, + "download": 2.0, + "extract": 4.0, + "parse": 8.0, + "embed": 16.0, + "db": 32.0, + "sweep": 64.0, + "other": 9.0, + } diff --git a/tests/unit/test_timing.py b/tests/unit/test_timing.py new file mode 100644 index 0000000..7c02a09 --- /dev/null +++ b/tests/unit/test_timing.py @@ -0,0 +1,86 @@ +"""Unit tests for indexer.timing: the ambient per-phase accumulator. + +Two properties the whole instrumentation rests on are pinned here, away from the +job pipeline that consumes them: ``record`` is a no-op when nobody is measuring +(so ``index_repo`` stays callable directly), and one thread's totals are +invisible to another (so fan-out cannot cross-attribute a phase). +""" + +from __future__ import annotations + +import threading + +import pytest + +from indexer.timing import PhaseTimer, current_timer, install_timer, record, reset_timer + + +@pytest.mark.unit +def test_timer_accumulates_and_is_a_noop_when_unset() -> None: + """Default-unset must be silent, not merely harmless. + + ``tests/integration/test_store.py`` calls ``index_repo`` directly with no + timer in context; the sweep hook in ``store.py`` runs there regardless, so + ``record`` with nothing installed has to do nothing AND not raise. + """ + assert current_timer() is None + record("sweep", 1.0) # must not raise + assert current_timer() is None + + timer = PhaseTimer() + assert timer.total("sweep") == 0.0 # an unrecorded phase reads as zero, not KeyError + + token = install_timer(timer) + try: + assert current_timer() is timer + record("sweep", 1.5) + record("sweep", 0.25) # accrues, never overwrites + record("parse", 2.0) + finally: + reset_timer(token) + + assert timer.total("sweep") == 1.75 + assert timer.total("parse") == 2.0 + assert timer.total("download") == 0.0 + + # The reset really uninstalled it: a later record goes nowhere. + assert current_timer() is None + record("sweep", 100.0) + assert timer.total("sweep") == 1.75 + + +@pytest.mark.unit +def test_timer_is_isolated_per_thread() -> None: + """Each thread starts with a fresh context, so timers cannot cross-attribute. + + This is what lets one ``PhaseTimer`` per branch be correct while four worker + threads index four repos concurrently. + """ + totals: dict[str, float] = {} + start = threading.Barrier(2) + + def worker(label: str, amount: float) -> None: + # A new thread inherits no context: nothing is installed here yet. + assert current_timer() is None + timer = PhaseTimer() + token = install_timer(timer) + try: + start.wait(timeout=5) # force real overlap, not sequential execution + record("db", amount) + start.wait(timeout=5) + totals[label] = timer.total("db") + finally: + reset_timer(token) + + threads = [ + threading.Thread(target=worker, args=("a", 3.0)), + threading.Thread(target=worker, args=("b", 7.0)), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + assert not thread.is_alive() + + assert totals == {"a": 3.0, "b": 7.0} + assert current_timer() is None # and nothing leaked back to the main thread From 17aeb4f16bc37e7b5ee8e1818401cc3f15112404 Mon Sep 17 00:00:00 2001 From: Tanner <84605639+IceRhymers@users.noreply.github.com> Date: Sat, 25 Jul 2026 08:18:34 -0700 Subject: [PATCH 3/9] indexer: file-level delta indexing keyed on (path, content_sha) (#104) Refs #104 --- app/config.py | 7 + app/db/models.py | 17 +- docs/runbooks/indexing-parallelism.md | 127 +++- docs/runbooks/semantic-enablement.md | 16 + indexer/job.py | 161 ++++- indexer/store.py | 352 +++++++++- tests/integration/test_store.py | 90 ++- tests/integration/test_store_chunk_writer.py | 125 +++- tests/integration/test_store_delta.py | 658 +++++++++++++++++++ tests/unit/test_job.py | 344 ++++++++++ tests/unit/test_store_delta.py | 523 +++++++++++++++ 11 files changed, 2379 insertions(+), 41 deletions(-) create mode 100644 tests/integration/test_store_delta.py create mode 100644 tests/unit/test_store_delta.py diff --git a/app/config.py b/app/config.py index e9900f8..e6bb63a 100644 --- a/app/config.py +++ b/app/config.py @@ -99,6 +99,13 @@ class Settings(BaseSettings): # this loud check could ever fire, which would defeat the point of having a ceiling. # A repo that legitimately exceeds this needs a temp-table staging path, not a bigger # buffer. + # + # Scope note (#104): under file-level delta indexing this cap is enforced against + # whatever ONE RUN embeds (changed/new + membership-only files), not a branch's whole + # corpus -- a branch can legitimately drift above this number between full reindexes + # (a semantics bump, or its first index), re-enforced in full at each of those. This is + # a deliberate, accepted trade-off (see indexer/job.py's module docstring and + # docs/runbooks/indexing-parallelism.md §4.1), not a bug; the constant is unchanged. semantic_max_chunks_per_repo: int = 8000 # Chunk size bound (tokens) fed to the embedding model. Distinct from MAX_FILE_BYTES, diff --git a/app/db/models.py b/app/db/models.py index 45c7e1e..af616db 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -45,7 +45,22 @@ ``indexer/symbols.py``, to ``indexer/parse.py``'s chunking, or to ``indexer/languages.py``'s extraction contract. A bump forces every repo to re-index once, because a repo's stored ``repos.index_semantics_version`` no -longer matches. The CI tripwire enforces the bump obligation. +longer matches. The CI tripwire enforces the bump obligation for those three +files. + +**The obligation extends past the tripwire's reach (#104).** Swapping the +embedding MODEL (``app/embed.py``) or changing ``SEMANTIC_EMBEDDING_DIM`` +(``app/config.py``) also requires a bump, but the tripwire does not watch +either file (``app/embed.py`` deliberately -- it would otherwise fire on +unrelated retry/batching edits) -- this is a reviewed convention, not a +machine-enforced one. Before file-level delta indexing (issue #104) a missed +bump here was self-limiting: the next HEAD move re-embedded a branch's whole +corpus regardless. Under delta indexing only a CHANGED file re-embeds, so a +missed bump now leaves every unchanged file's vectors silently stale forever +-- exactly the failure this version column exists to prevent. This is not a +new kind of case: version ``2`` was minted for precisely this reason (turning +semantic search on by default, so every already-indexed branch had to +re-index once for ``chunks`` to backfill). Migrations must never import this constant -- see ``app/alembic/versions/0002_index_semantics_version.py``. diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index fed089c..9de234b 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/docs/runbooks/indexing-parallelism.md @@ -154,6 +154,16 @@ you write keeps working. Read the largest field; that is the branch's bottleneck | `db` | per-file round trips | #105 (batched writes) | | any of the above, on **unchanged** content | redundant work | #104 (file-level delta indexing) | +`#104` narrows the **db** and **embed** costs to a branch's actual delta, not +its size — but it does NOT touch `parse`: extraction still runs on every +file every run (tree-sitter must produce a `FileExtraction` before +`index_repo` can classify it), so an all-unchanged branch on a large repo +still pays its full `download`+`extract`+`parse` cost. See `indexer.store`'s +`delta write set …` line (below) to tell "this branch is genuinely mostly-new" +from "this branch is mostly-unchanged but still parsing everything" — the +latter is exactly the case `#108` (process-pool extraction) or a future +extraction-skip step would address next. + Four fields need interpretation before you act on them: - **`resolve=0.00s` on a default branch is expected, not a bug.** That branch's @@ -184,6 +194,36 @@ Four fields need interpretation before you act on them: instrumentation merely makes visible for the first time — it is not a new regression. +### 2.2 The delta write set line (#104) + +Every `index_repo` call also emits one `indexer.store` INFO line, immediately +before the sweep, in **both** the gate-open and gate-closed cases — one format +string, no conditional fields, so it stays greppable either way: + +``` +INFO indexer.store [acme/widgets]: acme/widgets@main: delta write set 412/30214 files (unchanged=29790 membership=12, semantics gate open) +INFO indexer.store [acme/gadgets]: acme/gadgets@main: delta write set 812/812 files (unchanged=0 membership=0, semantics gate closed: stored v3 != v4) +``` + +`unchanged` files write nothing at all (no file upsert, no symbol/edge +delete-reinsert, no chunk write) and are never re-embedded. `membership` files +are already stored under another branch and only need their `branches` array +unioned in, plus a chunk write if semantic is on (see §4's accepted +regressions). The leading fraction is `(changed/new) / (total seen)`. The gate +is per-BRANCH: it opens only once that branch's own `repo_branches` stamp is +at the current `INDEX_SEMANTICS_VERSION` — a branch's first run, or any run +after a semantics bump, always shows `semantics gate closed`. + +``` +grep 'delta write set' run.log # one line per index_repo call +``` + +A branch stuck at a low `unchanged=` fraction run after run either genuinely +churns every run (nothing to fix) or has drifted out of delta eligibility — +check its `repo_branches.index_semantics_version` against the current +`INDEX_SEMANTICS_VERSION` and whether a sibling branch is stale (§4's +provenance gate). + --- ## 3. The three limits, and why raising concurrency is a bad trade @@ -265,21 +305,83 @@ There is deliberately **no `--force_reindex` flag.** Forcing a re-index means clearing the provenance stamp, after which the normal skip logic re-indexes the affected repos on the next scheduled or manual run. +**The stamp the skip seam actually reads is `repo_branches`, not `repos`.** +`indexer/job.py`'s `_read_stamps` selects +`RepoBranch.last_indexed_commit, RepoBranch.index_semantics_version` — the +`repos` table's `index_semantics_version` column is a deprecated legacy stamp +that no decision anywhere reads (`app/db/models.py` documents it write-only). +An `UPDATE repos SET index_semantics_version = NULL` is therefore a **no-op** +against the skip seam: the branch will look untouched and re-index on its own +next scheduled cycle, not immediately, and the operator following an older +version of this runbook would see nothing happen. + ```sql -- everything -UPDATE repos SET index_semantics_version = NULL; +UPDATE repo_branches SET index_semantics_version = NULL; + +-- one repo, every branch +UPDATE repo_branches SET index_semantics_version = NULL + WHERE repo_id = (SELECT id FROM repos WHERE name = 'acme/widgets'); --- one repo -UPDATE repos SET index_semantics_version = NULL WHERE name = 'acme/widgets'; +-- one repo, one branch +UPDATE repo_branches SET index_semantics_version = NULL + WHERE repo_id = (SELECT id FROM repos WHERE name = 'acme/widgets') + AND branch = 'main'; ``` Then run the job (`make index TARGET=` or `databricks bundle run code_search_index -t `). +### 4.1 File-level delta indexing (#104): what changes about this remedy + +Once a branch's `repo_branches.index_semantics_version` matches the current +`INDEX_SEMANTICS_VERSION`, `index_repo` skips rewriting any file whose +`(path, content_sha)` it already has stored for that branch — see +`indexer.store`'s module docstring for the full classification and the +correctness proof. Two consequences change what "clear the stamp" actually +buys you: + +**A degraded branch no longer self-heals on its own.** Before #104, ANY +re-index rewrote the whole branch, so a branch whose semantic precompute +failed (a chunk-cap breach, an embedder outage) caught its chunks up +automatically on the next successful run. Under delta indexing, only +*changed* files get re-embedded — a branch that never changes again carries +that gap **forever** unless you clear its stamp. `indexer.job` emits one +run-completion WARNING naming every branch that finished this way: + +``` +WARNING indexer.job [-]: 2 branch(es) finished with degraded semantic coverage this run (chunk precompute failed; core index is current, chunks are not, and delta indexing will NOT catch them up on their own -- clear their repo_branches.index_semantics_version stamp to force a full re-embed, see docs/runbooks/indexing-parallelism.md §4): acme/big-repo@main, acme/other@release +``` + +Grep for it (`grep 'degraded semantic coverage' run.log`) and clear the named +branches' stamps with the one-branch form above once the underlying cause +(chunk cap, embedder outage) is resolved. + +**The provenance gate can force a full re-index you did not ask for.** A +branch only takes the cheaper "membership-only" path (acquiring content a +*sibling* branch already stored, e.g. two branches sharing most of a +monorepo) when **every** `repo_branches` row for that repo is at the current +semantics version. If you clear one branch's stamp and leave siblings +untouched, that is fine — but a repo with one branch stuck at an old version +for any other reason (a persistently failing branch) will force every OTHER +branch of that repo through the full write path for any file it shares with +the stuck one, even though those branches are otherwise fully caught up. The +`delta write set …` line's `membership=` count going to zero across a whole +repo, with `unchanged=` still high, is the symptom — check for a sibling +branch stuck at a stale `index_semantics_version` before assuming something +is broken. + +**`semantic_max_chunks_per_repo` is enforced per RUN, not per branch's whole +corpus.** The cap is evaluated over whatever `_precompute_chunk_writer` +embeds, which under delta indexing is only the changed/new/membership-only +files. A branch can drift above the nominal cap between full reindexes (a +semantics bump, or its first index) — re-enforced in full at each of those. +Not a bug; see `app/config.py`'s `semantic_max_chunks_per_repo` comment. + ### Who can run this — read before you need it -`UPDATE` on `repos` is held by **the identity that deployed the schema**, which -owns the tables. Concretely: +`UPDATE` on `repo_branches` (and `repos`) is held by **the identity that deployed +the schema**, which owns every table, `repo_branches` included. Concretely: - **dev:** the developer who ran `make migrate` / `scripts/deploy.sh`. Table ownership carries `UPDATE` implicitly; no explicit grant was ever issued for @@ -305,6 +407,21 @@ If you change **what** gets extracted — `indexer/symbols.py`, `indexer/parse.py`, `indexer/languages.py` — you **must** bump `INDEX_SEMANTICS_VERSION` in `app/db/models.py`. +**The same obligation now extends past the tripwire's watched files (#104).** +`indexer/parse.py`'s chunker is already a watched path, so a change to +`iter_chunks` still fires the tripwire. Swapping the embedding MODEL +(`app/embed.py`) or changing `SEMANTIC_EMBEDDING_DIM` (`app/config.py`) +without bumping `INDEX_SEMANTICS_VERSION` is NOT caught by the tripwire (both +are deliberately unwatched — `app/embed.py` would otherwise fire on +unrelated retry/batching edits and a noisy tripwire gets disabled) and now +leaves every UNCHANGED file's vectors permanently stale under file-level +delta indexing — before #104 the next HEAD move re-embedded everything +anyway, so a missed bump here was self-limiting; it no longer is. This is not +a new pattern: `INDEX_SEMANTICS_VERSION` version `2` was minted for exactly +this reason (turning semantic search on by default so `chunks` backfills). +Treat this as a reviewed convention, the same posture `indexer.store`'s +module docstring takes for `lang`/`size` re-derivation. + Without a bump, every already-indexed repo keeps serving output from the *old* extractor and never re-indexes, because its stored stamp still matches HEAD. The failure is silent and open-ended: the index looks perfectly current. diff --git a/docs/runbooks/semantic-enablement.md b/docs/runbooks/semantic-enablement.md index ee195e0..51ba874 100644 --- a/docs/runbooks/semantic-enablement.md +++ b/docs/runbooks/semantic-enablement.md @@ -85,6 +85,22 @@ which makes the job a true semantic no-op (no embedder built, no chunking, the environment. Precedence for the job is `config.yaml > CODE_SEARCH_* env > default`, so `semantic.enabled: false` wins even if the env says enabled. +**Re-enabling the job's semantic flag does not backfill on its own (#104).** This +runbook's §2 above promises that an `INDEX_SEMANTICS_VERSION` bump "forces every +already-indexed branch to re-index once … which backfills `chunks`" — true for a +version bump, but **not** for flipping `semantic.enabled` back to `true` after a +period disabled. Under file-level delta indexing a branch whose stamp is already at +the current `INDEX_SEMANTICS_VERSION` classifies every unchanged file as unchanged +and skips embedding it, with no awareness that this run is the first to have an +embedder at all. A branch that never changes again after re-enabling never gets +chunks. Clear that branch's stamp to force the backfill (same remedy as a degraded +branch — see `docs/runbooks/indexing-parallelism.md` §4.1): + +```sql +UPDATE repo_branches SET index_semantics_version = NULL + WHERE repo_id = (SELECT id FROM repos WHERE name = 'acme/widgets'); +``` + **All three surfaces, not just one:** the flag must be off on the MCP app, the webui app (both via env), **and** the indexer job (via `config.yaml`) — each has its own config source. The webui SPA's Semantic tab is driven entirely by diff --git a/indexer/job.py b/indexer/job.py index 3a0f1d8..704209f 100644 --- a/indexer/job.py +++ b/indexer/job.py @@ -79,6 +79,24 @@ loop as symbols. Flag-off: no chunking, no embedder, no import of ``app.embed``'s lazy ``databricks-sdk`` dependency. +File-level delta indexing (issue #104) skips re-embedding a file this branch +already carries unchanged, at ``(path, content_sha)``. Before building the +embed list, this module calls the injected ``shas_fn`` (default +:func:`indexer.store.read_indexed_shas`) -- but ONLY when this branch's stored +``index_semantics_version`` already matches +:data:`app.db.models.INDEX_SEMANTICS_VERSION`, the same gate +``indexer.store.index_repo`` applies authoritatively inside its own +transaction -- on a short-lived connection that closes before the embedder is +called. ``index_repo``'s per-file classification (unchanged / membership-only / +changed-new) is the one that actually decides what gets written; this module's +copy is advisory only, used to decide what needs a vector. Both reads can only +diverge if a second writer touched this repo between them, which the +single-writer-per-repo invariant above forbids. A branch whose semantic +precompute fails keeps whatever chunk coverage it already had -- see +:func:`_precompute_chunk_writer` and ``indexer/store.py``'s module docstring +for why that is no longer self-healing on the next run under delta indexing, +and the aggregate run-completion WARNING this module emits for it. + Every INDEXED branch also emits one ``phase timing`` line accounting for its whole wall clock -- resolve / download / extract / parse / embed / db / sweep, plus an ``other`` residual so no time is silently unattributed. The fields are fixed and @@ -130,15 +148,18 @@ resolve_branch_head, resolve_ref, ) +from indexer.hashing import content_sha from indexer.languages import Chunk, FileExtraction, IndexCounts, ParsedFile from indexer.parse import iter_chunks, iter_source_files from indexer.repo_config import RepoConfig, effective_workers, load_config, normalize_repo from indexer.resolve import MAX_REPOS, RepoEntry, resolve_repos from indexer.store import ( ChunkWriter, + ContentShaSets, ReconcileCounts, StaleIndexError, index_repo, + read_indexed_shas, reconcile_removed_repos, reconcile_retired_branches, ) @@ -166,11 +187,20 @@ class BranchOutcome: ``StaleIndexError`` maps to ``"conflict"`` and any other exception to ``"failed"``, caught INSIDE the per-branch loop so one branch's failure never stops its repo's other branches from being attempted. + + ``semantic_degraded`` is ``True`` only for an ``"indexed"`` outcome whose + chunk precompute raised (a chunk-cap breach or an embedder failure) -- + never for semantic-off, and never for a build-time embedder misconfiguration + (that degrades ``embed_fn`` to ``None`` before any branch starts, so no + per-branch precompute is ever attempted for it -- see ``_index_one_branch``). + ``run()`` aggregates every branch with this flag set into one + run-completion WARNING. """ branch: str status: Literal["indexed", "skipped", "conflict", "failed"] counts: IndexCounts | None = None + semantic_degraded: bool = False @dataclass(frozen=True) @@ -276,6 +306,7 @@ def run( max_repos: int = MAX_REPOS, reconcile_retired_fn: Callable[..., ReconcileCounts] = reconcile_retired_branches, reconcile_removed_fn: Callable[..., list[str]] = reconcile_removed_repos, + shas_fn: Callable[..., ContentShaSets] = read_indexed_shas, ) -> int: """Index every configured repo and return a process exit code (0 = all ok). @@ -414,6 +445,7 @@ def run( # -- and how a shrunken disk becomes visible before it becomes an outage. ok = skipped = conflicts = failures = 0 repo_outcomes: list[RepoOutcome] = [] + degraded_branches: list[str] = [] reconciliation_attempted = False reconciliation_failed = False reconcile_skip_reason = "" @@ -451,6 +483,7 @@ def run( cfg=cfg, embed_fn=embed_fn, stamps=stamps, + shas_fn=shas_fn, ): entry for entry in entries } @@ -479,6 +512,8 @@ def run( else: ok += 1 assert outcome.counts is not None + if outcome.semantic_degraded: + degraded_branches.append(f"{entry.name}@{outcome.branch}") logger.info( "indexed %s@%s: files=%d symbols=%d edges=%d swept=%d", entry.name, @@ -548,6 +583,28 @@ def run( # the WARNING logged at the conflict site is the record. It is NOT that the # work was redundant. + # One aggregate, greppable WARNING for every branch that finished "indexed" + # with degraded semantic coverage this run (its precompute failed -- a + # chunk-cap breach or an embedder outage -- so its core index is current but + # its chunks are not). Under file-level delta indexing this gap is NOT + # self-healing on the next run (see indexer/store.py's module docstring and + # docs/runbooks/indexing-parallelism.md §4): only a changed file re-embeds, + # so a branch that never changes again would carry stale/missing chunks + # forever unless an operator clears its semantics stamp. The per-branch + # WARNING already logged at the precompute site is easy to miss in a large + # run's log; this line exists so the condition is greppable after the fact. + # Deliberately does NOT fail the run -- see the per-branch warning site for + # why this is an additive-layer failure, not a core-index one. + if degraded_branches: + logger.warning( + "%d branch(es) finished with degraded semantic coverage this run (chunk precompute " + "failed; core index is current, chunks are not, and delta indexing will NOT catch " + "them up on their own -- clear their repo_branches.index_semantics_version stamp to " + "force a full re-embed, see docs/runbooks/indexing-parallelism.md §4): %s", + len(degraded_branches), + ", ".join(sorted(degraded_branches)), + ) + # Exactly one reconciliation summary line, always after "indexing complete". # The failure/withheld path already logged its own ERROR incident line # inside _reconcile -- this branch must not double-log it. @@ -803,8 +860,21 @@ def _precompute_chunk_writer( embedder itself. Raises ``ValueError`` if the repo's total chunk count exceeds ``max_chunks_per_repo`` (the documented hard ceiling, not a streaming bound -- see ``app.config.semantic_max_chunks_per_repo``). + + ``files`` is whatever the caller decided needs vectors -- under file-level + delta indexing (``_index_one_branch``) that is every file the advisory + ``shas_fn`` read did NOT classify as unchanged, not necessarily every parsed + file in the branch. The closure below therefore closes over ``covered = + set(per_file)`` -- every path THIS call embedded, including zero-chunk files + -- and refuses to write chunks for any other path. This is defence-in-depth + for a path the single-writer-per-repo invariant says is unreachable: ` + `write_chunks`` deletes a file's chunk rows before inserting, so calling it + for a path this precompute never embedded would silently delete that file's + chunks rather than merely leave them stale. See ``indexer/store.py``'s + ``_union_membership`` for the authoritative-side analogue of this guard. """ per_file: dict[str, list[Chunk]] = {pf.path: list(iter_chunks(pf)) for pf in files} + covered = set(per_file) total = sum(len(chunks) for chunks in per_file.values()) if total > max_chunks_per_repo: raise ValueError( @@ -834,6 +904,17 @@ def _precompute_chunk_writer( i += len(chunks) def chunk_writer(conn: Any, repo_id: int, file_id: int, pf: ParsedFile) -> None: + if pf.path not in covered: + # Unreachable while the single-writer invariant holds (see the + # docstring): index_repo only ever calls chunk_writer for a file this + # same precompute either embedded or classified membership-only (and + # index_repo's own _union_membership guards that case separately). + # Warn-and-skip rather than raise, matching this module's established + # additive-layer posture for the semantic path. + logger.warning( + "no precomputed chunks for %s; leaving its chunk rows untouched", pf.path + ) + return write_chunks(conn, file_id=file_id, chunks=by_path.get(pf.path, [])) return chunk_writer @@ -848,6 +929,7 @@ def _index_one( cfg: Settings, embed_fn: EmbedFn | None, stamps: dict[tuple[str, str], tuple[str | None, int | None]], + shas_fn: Callable[..., ContentShaSets], ) -> RepoOutcome: """Run the full fetch -> parse -> symbols -> store pipeline for every branch of one repo. @@ -873,6 +955,7 @@ def _index_one( cfg=cfg, embed_fn=embed_fn, stamps=stamps, + shas_fn=shas_fn, ) finally: _repo_ctx.reset(token) @@ -888,6 +971,7 @@ def _index_one_inner( cfg: Settings, embed_fn: EmbedFn | None, stamps: dict[tuple[str, str], tuple[str | None, int | None]], + shas_fn: Callable[..., ContentShaSets], ) -> RepoOutcome: """The body of :func:`_index_one`, run with the repo log context already set. @@ -948,6 +1032,7 @@ def _index_one_inner( stamps=stamps, started=started, max_chunks_per_repo=max_chunks_per_repo, + shas_fn=shas_fn, ) for branch in resolution.branches ] @@ -1015,6 +1100,7 @@ def _index_one_branch( stamps: dict[tuple[str, str], tuple[str | None, int | None]], started: float, max_chunks_per_repo: int, + shas_fn: Callable[..., ContentShaSets], ) -> BranchOutcome: """Fetch, parse, and store ONE branch. Never raises -- every failure is classified. @@ -1023,6 +1109,22 @@ def _index_one_branch( A stored ``None`` version means the provenance of the stored index is unknown, so the branch is always re-indexed. + ``shas_fn`` is the ADVISORY copy of :func:`indexer.store.read_repo_content_shas` + (see that module's docstring for the authoritative one). It is called ONLY + when this branch's stored ``index_semantics_version`` already matches + :data:`app.db.models.INDEX_SEMANTICS_VERSION` -- the same gate ``index_repo`` + applies inside its transaction, from data already in hand here -- and ONLY on + a separate, short-lived ``engine.connect()`` that closes before embedding + starts, never on a connection held across the embedder's network I/O. Its + result narrows the file list handed to ``_precompute_chunk_writer`` to every + file NOT already carried by this branch (changed/new *and* membership-only -- + ``index_repo`` may reclassify a membership-only file as changed/new inside its + own transaction if the provenance gate fails there, so it must already have a + vector to attach). The two reads can only disagree if a second writer touched + this repo between them, which the single-writer-per-repo invariant (see the + module docstring) forbids; see ``indexer/store.py``'s ``_union_membership`` + for the defence-in-depth guard on the authoritative side. + ``max_chunks_per_repo`` is the caller's (``_index_one_inner``'s) already-resolved effective cap -- this repo's ``semantic_max_chunks_per_repo`` override if one matched, else ``cfg.semantic_max_chunks_per_repo``. Taken as a parameter rather @@ -1083,6 +1185,7 @@ def _index_one_branch( timer.add("extract", timer.clock() - t0) chunk_writer: ChunkWriter | None = None + precompute_failed = False if cfg.semantic_enabled and embed_fn is not None: # Chunking/embedding needs the full file list up front -- unlike # the lazy items generator below, it cannot stream through @@ -1100,17 +1203,56 @@ def _index_one_branch( # In a `finally`, unlike every other phase wrap: the degrade path # below still burned this time (a downed embedder can burn a lot - # of it before it gives up) and must still be reported. + # of it before it gives up) and must still be reported. The + # advisory shas_fn read is charged here too, deliberately NOT as + # its own timed phase: it exists solely to decide what this block + # embeds, and #103's `phase timing` line is pinned exhaustive + # (nine fixed fields, tests/unit/test_job.py) -- adding a tenth + # field is out of this change's scope. t0 = timer.clock() try: - chunk_writer = _precompute_chunk_writer(files, embed_fn, max_chunks_per_repo) + files_to_embed = files + if stamps.get((name.casefold(), branch), (None, None))[1] == ( + INDEX_SEMANTICS_VERSION + ): + # Delta gate open (same test index_repo will apply + # authoritatively, from data already in hand): narrow to + # every file this branch does not already carry. A short- + # lived connection, closed before embedding starts -- + # never held across the embedder's network I/O. + # + # `_present` is discarded -- this module only ever needs + # `carried` (job.py cannot replicate index_repo's + # provenance-gate check anyway, and doesn't need to: any + # not-carried file gets embedded here regardless of + # whether index_repo later classifies it membership-only + # or changed/new). shas_fn still computes and returns the + # full-repo `present` set -- read_repo_content_shas' + # signature is deliberately ONE shared query pair with + # index_repo's authoritative read (see its docstring), so + # this module pays for a second full-repo Index Only Scan + # it doesn't use rather than forking the query. That cost + # rides inside `embed=` on the phase timing line (see the + # comment above), not broken out separately. + with engine.connect() as shas_conn: + carried, _present = shas_fn(shas_conn, name=name, branch=branch) + files_to_embed = [ + pf for pf in files if (pf.path, content_sha(pf.content)) not in carried + ] + chunk_writer = _precompute_chunk_writer( + files_to_embed, embed_fn, max_chunks_per_repo + ) except Exception: # The semantic layer is ADDITIVE: a chunk-ceiling breach, a downed embedder, - # or a dim/count mismatch must not cost this branch its core index. Letting - # it propagate would skip files/symbols AND the mark-and-sweep, silently - # leaving the branch stale -- worse than stale chunks. Chunks catch up on - # the next successful run; the failure is logged with a traceback, never - # swallowed silently. + # a dim/count mismatch, or a failure reading the advisory shas_fn projection + # must not cost this branch its core index. Letting it propagate would skip + # files/symbols AND the mark-and-sweep, silently leaving the branch stale -- + # worse than stale chunks. Under file-level delta indexing this is NOT + # self-healing the way it was before: only a changed file re-embeds, so a + # branch that never changes again carries this gap forever unless an + # operator clears its semantics stamp (see run()'s aggregate WARNING and + # docs/runbooks/indexing-parallelism.md §4). The failure is logged with a + # traceback here too, never swallowed silently. logger.warning( "semantic precompute failed for %s@%s; indexing core corpus without chunks", name, @@ -1118,6 +1260,7 @@ def _index_one_branch( exc_info=True, ) chunk_writer = None + precompute_failed = True finally: timer.add("embed", timer.clock() - t0) items = ((pf, extract_file(pf)) for pf in files) @@ -1198,7 +1341,9 @@ def _index_one_branch( ) except Exception: logger.warning("phase timing unavailable for %s@%s", name, branch, exc_info=True) - return BranchOutcome(branch=branch, status="indexed", counts=counts) + return BranchOutcome( + branch=branch, status="indexed", counts=counts, semantic_degraded=precompute_failed + ) except StaleIndexError as exc: # The repo_branches row for THIS branch changed under this worker, so # its whole transaction rolled back and THIS BRANCH IS NOT INDEXED. diff --git a/indexer/store.py b/indexer/store.py index 84e110a..17f8c63 100644 --- a/indexer/store.py +++ b/indexer/store.py @@ -69,6 +69,81 @@ class ReconcileCounts: # computed -- this seam never calls an embedder itself. ChunkWriter = Callable[[Connection, int, int, ParsedFile], None] +# The two projection reads behind the file-level delta path. Returned as +# ``(carried, present)`` -- see read_repo_content_shas. +ContentShaSets = tuple[set[tuple[str, str]], set[tuple[str, str]]] + + +def read_repo_content_shas(conn: Connection, *, repo_id: int, branch: str) -> ContentShaSets: + """Project one repo's stored ``(path, content_sha)`` pairs, twice. + + Returns ``(carried, present)``: + + * ``carried`` -- the pairs on rows whose ``branches`` array already contains + ``branch``. A parsed file in this set is UNCHANGED for this branch. + * ``present`` -- every pair stored for this repo, on any branch. A parsed + file in ``present - carried`` already exists as a row this branch does not + yet carry, i.e. the membership-only class. + + **The projection is ``path, content_sha`` and nothing else, deliberately.** + Two expensive mistakes are available here and both must stay closed: + + * selecting ``content`` pulls the entire corpus into the worker and defeats + the whole point of the delta path; + * selecting ``branches`` forces a heap fetch per row on a table whose rows + carry that content, so the second read stops being an Index Only Scan. + + **They are also two separate statements on purpose.** Do NOT collapse them + into ``SELECT path, content_sha, branches @> ... AS carries FROM files WHERE + repo_id = :id``: that drops the branch predicate entirely (so + ``ix_files_branches_gin`` is never consulted) *and* projects ``branches``. + The containment form ``branches @> ARRAY[:branch]`` is what the GIN index can + serve; ``_sweep_membership``'s ``:branch = ANY(branches)`` cannot be. + + **Keyed on ``(path, content_sha)``, NEVER on path alone.** + ``uq_files_repo_path_sha`` permits several rows for one path with different + content (the divergent-branch case), so a path-keyed dict silently drops rows + and which one survives depends on row order. + """ + carried = { + (row.path, row.content_sha) + for row in conn.execute( + text( + "SELECT path, content_sha FROM files " + "WHERE repo_id = :repo_id AND branches @> CAST(:branch_arr AS text[])" + ), + {"repo_id": repo_id, "branch_arr": [branch]}, + ) + } + present = { + (row.path, row.content_sha) + for row in conn.execute( + text("SELECT path, content_sha FROM files WHERE repo_id = :repo_id"), + {"repo_id": repo_id}, + ) + } + return carried, present + + +def read_indexed_shas(conn: Connection, *, name: str, branch: str) -> ContentShaSets: + """Name-keyed wrapper around :func:`read_repo_content_shas` for ``indexer.job``. + + Resolves ``repos.name -> id`` itself and returns two empty sets for a repo + that has never been indexed (which degrades to "everything is changed/new" -- + safe in the correct direction). This is the ADVISORY copy of the read: the + authoritative one runs inside ``index_repo``'s transaction. Both go through + the same helper so there is exactly one pair of queries and one keying rule. + + Called on its own short-lived connection that is closed BEFORE embedding + starts -- never on a connection held across network I/O. + """ + repo_id = conn.execute( + text("SELECT id FROM repos WHERE name = :name"), {"name": name} + ).scalar_one_or_none() + if repo_id is None: + return set(), set() + return read_repo_content_shas(conn, repo_id=int(repo_id), branch=branch) + def index_repo( conn: Connection, @@ -95,14 +170,30 @@ def index_repo( 2. Upsert/read the ``repo_branches`` row for ``(repo_id, branch)`` under its row lock, capturing ``(baseline_commit, baseline_version)`` -- the CAS baseline for step 5, mirroring the same ``RETURNING`` trick as step 1. - 3. Per file: an array-union upsert on ``uq_files_repo_path_sha`` -- a file - whose content already exists under another branch gets THIS branch - unioned into its ``branches`` array (one row, shared content); a file - whose content differs from every existing version gets its own row. Then - delete-and-reinsert its ``symbols`` and ``reference_edges`` (neither has a - natural key), then call ``chunk_writer`` (if given) so chunk writes - commit/roll back with the rest of that file's row. Each processed file's - ``(path, content_sha)`` is collected into this branch's seen-set. + 3. Two projection reads (:func:`read_repo_content_shas`), issued ONLY when + the delta gate is open -- see below -- then, per file, a three-way + classification on ``(pf.path, content_sha(pf.content))``: + + * **unchanged** (the pair is already on a row carrying this branch): no + statement at all. + * **membership-only** (the pair is stored for this repo but on a row this + branch does not carry, AND statement 4 proves every branch of this repo + is at the current semantics version): no symbol/edge work; the whole + class is unioned in by ONE batched ``UPDATE ... RETURNING`` after the + loop, which also supplies the ``file_id`` for its ``chunk_writer`` call + (see :func:`_union_membership`). + * **changed/new** (everything else): an array-union upsert on + ``uq_files_repo_path_sha`` -- a file whose content already exists under + another branch gets THIS branch unioned into its ``branches`` array (one + row, shared content); a file whose content differs from every existing + version gets its own row. Then delete-and-reinsert its ``symbols`` and + ``reference_edges`` (neither has a natural key), then call + ``chunk_writer`` (if given) so chunk writes commit/roll back with the + rest of that file's row. + + **Every** parsed file -- classified or written -- is collected into this + branch's seen-set, so step 4 and its empty-seen-set guard are correct by + construction and untouched by the delta path. 4. Membership sweep, keyed on THIS branch's seen-set (never on ``commit``, which is ambiguous under dedup): strip ``branch`` from any row's ``branches`` array that is not in the seen-set, then delete any row left @@ -112,17 +203,66 @@ def index_repo( set is empty** -- an empty seen-set would otherwise strip ``branch`` from every row in the repo; conservatively skipping is safer than wiping. 5. CAS-stamp the ``repo_branches`` row for ``(repo_id, branch)`` against the - step-2 baseline (raises :class:`StaleIndexError` on mismatch). + step-2 baseline (raises :class:`StaleIndexError` on mismatch). A run whose + seen-set was EMPTY advances ``last_indexed_commit`` but leaves + ``index_semantics_version`` at the step-2 baseline -- it wrote nothing, so + it indexed nothing at the current semantics version. See + :func:`_stamp_repo_branch`. ``items`` may be a lazy generator; it is consumed inside the open transaction so memory stays bounded. ``chunk_writer`` defaults to ``None``, which makes this byte-identical to the core (semantic-off) path; when given, it must write PRECOMPUTED chunks -- embeddings are computed outside this transaction, so no network call ever happens here. + + **The delta gate.** The classification above is taken only when + ``baseline_version == INDEX_SEMANTICS_VERSION`` -- statement 2's ``RETURNING`` + value, already in hand, no extra query. A ``NULL`` or older version means + every file takes the full write path and statements 3a/3b are never issued. + + Why the unchanged path is sound, inductively: + + * A row carrying this branch was necessarily written by this branch's last + COMPLETED run, and that run ran at ``baseline_version == + INDEX_SEMANTICS_VERSION``. In it the row was either written full-path (so + it is current), or skipped as unchanged (current, by induction). + * ... or acquired membership-only, which + :func:`_repo_is_wholly_at_current_version` only permits when every branch + of the repo is at the current version (so every surviving row of the repo + was last written at it). + * Base case: the first run after ANY version transition has + ``baseline_version != INDEX_SEMANTICS_VERSION``, so it is full-path for + every parsed file. A zero-parse run cannot manufacture a spurious base + case -- it does not advance the version stamp (see + :func:`_stamp_repo_branch`). + + Two columns are deliberately NOT re-derived for a skipped file: + + * ``lang`` and ``size`` are pure functions of ``(path, content)`` via + ``indexer/parse.py`` + ``indexer/languages.py``, both watched by + ``tests/unit/test_semantics_version_tripwire.py`` -- so a change to either + derivation is MEANT to force a version bump, which closes this gate. That + tripwire is a local-developer guard rather than a CI one, so treat this as + a strong convention backed by review, not a machine-enforced invariant. + * ``files.commit`` goes staler. No production read path exists (every + ``commit:`` filter resolves from ``repo_branches.last_indexed_commit``, + and the column is documented write-only and ambiguous under dedup in + ``app/db/models.py``). This makes it staler; it makes nothing wrong. + + The breakdown is reported on one INFO line per call, immediately before the + sweep, in both the gate-open and gate-closed cases:: + + acme/widgets@main: delta write set 412/30214 files (unchanged=29790 membership=12, + semantics gate open) + + ``IndexCounts`` is unchanged: ``files`` still counts files SEEN this run, and + ``symbols``/``edges`` still count rows actually inserted -- so they + legitimately read ``0`` on an all-unchanged run. That is the correct signal. """ file_count = 0 symbol_count = 0 edge_count = 0 + unchanged_count = 0 seen_paths: list[str] = [] seen_shas: list[str] = [] @@ -159,8 +299,47 @@ def index_repo( ) baseline_commit, baseline_version = conn.execute(branch_stmt).one() + # Statements 3a/3b/4, issued only behind the delta gate. Everything the + # classification below needs is now in hand; nothing else is read. + delta_on = baseline_version == INDEX_SEMANTICS_VERSION + carried: set[tuple[str, str]] = set() + present: set[tuple[str, str]] = set() + membership_ok = False + if delta_on: + carried, present = read_repo_content_shas(conn, repo_id=repo_id, branch=branch) + membership_ok = _repo_is_wholly_at_current_version(conn, repo_id=repo_id) + + # (pf, content_sha) for each membership-only file, held until the batched + # UPDATE below can hand back their file ids. A bounded exception to the + # "items stream through the transaction" rule: membership-only is the + # rare class (a branch ACQUIRING content another branch already stored), + # not the steady state, and chunk_writer's seam takes the ParsedFile. + membership: list[tuple[ParsedFile, str]] = [] + for pf, ex in items: sha = content_sha(pf.content) + # Seen-set membership is recorded for EVERY parsed file, whatever its + # class -- that is what keeps the sweep (and its empty-seen-set + # guard) correct without any delta awareness of its own. + file_count += 1 + seen_paths.append(pf.path) + seen_shas.append(sha) + + if delta_on and (pf.path, sha) in carried: + # Unchanged: this exact content is already stored on a row this + # branch already carries. No file upsert, no symbol/edge + # delete-reinsert, no chunk_writer call. + unchanged_count += 1 + continue + + if delta_on and membership_ok and (pf.path, sha) in present: + # Membership-only: the row exists (written by another branch) but + # does not carry this branch yet. Statement 4 has proven every + # branch of this repo is at the current semantics version, so its + # symbols/edges are current and only the array union is owed. + membership.append((pf, sha)) + continue + file_stmt = ( pg_insert(File) .values( @@ -194,9 +373,6 @@ def index_repo( .returning(File.id) ) file_id = conn.execute(file_stmt).scalar_one() - file_count += 1 - seen_paths.append(pf.path) - seen_shas.append(sha) conn.execute(delete(Symbol).where(Symbol.file_id == file_id)) if ex.symbols: @@ -243,6 +419,38 @@ def index_repo( if chunk_writer is not None: chunk_writer(conn, repo_id, file_id, pf) + # ONE statement for the whole membership-only class, skipped entirely + # when that class is empty (rather than issued as a no-op) so the + # statement inventory stays stable and greppable. + if membership: + _union_membership( + conn, + repo_id=repo_id, + branch=branch, + membership=membership, + chunk_writer=chunk_writer, + ) + + # One INFO line per index_repo call, immediately before the sweep, in + # BOTH the gate-open and gate-closed cases -- one format string, no + # conditional fields, so the line is always present and always greppable. + # The reason tail is the only part that varies. IndexCounts is + # deliberately NOT extended to carry this: it is a frozen dataclass + # compared by value in existing assertions, and `files` keeps meaning + # "files seen this run" (the seen-set size). + logger.info( + "%s@%s: delta write set %d/%d files (unchanged=%d membership=%d, %s)", + name, + branch, + file_count - unchanged_count - len(membership), + file_count, + unchanged_count, + len(membership), + "semantics gate open" + if delta_on + else f"semantics gate closed: stored v{baseline_version} != v{INDEX_SEMANTICS_VERSION}", + ) + # Timed into indexer.job's ambient per-branch PhaseTimer, if one is # installed -- a no-op otherwise, so a direct index_repo call (tests, # scripts) is unaffected. Deliberately NOT a return value: IndexCounts is @@ -269,11 +477,98 @@ def index_repo( head_sha=head_sha, baseline_commit=baseline_commit, baseline_version=baseline_version, + seen_any=bool(seen_paths), ) return IndexCounts(files=file_count, symbols=symbol_count, swept=swept, edges=edge_count) +def _repo_is_wholly_at_current_version(conn: Connection, *, repo_id: int) -> bool: + """Statement 4: is EVERY ``repo_branches`` row for this repo at the current version? + + The provenance gate the membership-only class depends on, and the hole + ``(path, content_sha)`` alone does not close. Counter-example it exists for: + branch ``b`` is stamped at the current version (delta on); sibling branch + ``a`` was written at an OLDER version and has not re-indexed since. ``b``'s + HEAD moves and acquires a file whose exact ``(path, content)`` already exists + as ``a``'s stale-version row. Taking the membership-only path would skip the + symbol/edge rewrite, so ``b`` would serve old-extractor symbols under a + current-version stamp -- silently, and exactly the failure + ``INDEX_SEMANTICS_VERSION`` exists to prevent. + + Given this gate, every surviving ``files`` row of the repo was last written + at the current version: every row carries at least one branch (both sweep + sites delete rows at ``cardinality(branches) = 0``), and every branch string + on a ``branches`` array has a ``repo_branches`` row (``index_repo`` writes + statement 2 before any file row for that branch, and + ``reconcile_retired_branches`` deletes both in one transaction). Both + directions are load-bearing and both are pinned by tests. + """ + return bool( + conn.execute( + text( + "SELECT NOT EXISTS (SELECT 1 FROM repo_branches " + "WHERE repo_id = :repo_id " + "AND index_semantics_version IS DISTINCT FROM :version)" + ), + {"repo_id": repo_id, "version": INDEX_SEMANTICS_VERSION}, + ).scalar_one() + ) + + +def _union_membership( + conn: Connection, + *, + repo_id: int, + branch: str, + membership: list[tuple[ParsedFile, str]], + chunk_writer: ChunkWriter | None, +) -> None: + """Union ``branch`` into every membership-only row in ONE statement, then write their chunks. + + ``array_agg(DISTINCT ...)`` rather than ``||`` alone so the stored array + stays sorted-distinct, matching ``index_repo``'s per-file upsert idiom -- + existing assertions compare ``branches`` by value. + + ``RETURNING id, path, content_sha`` supplies each row's ``file_id`` without a + second lookup, which is what makes the ``chunk_writer`` call below possible. + **Membership-only DOES write chunks** even though it writes no symbols or + edges: the acquired row may legitimately have zero chunk rows (the branch + that first wrote it ran semantic-off, or its precompute failed), and skipping + the write would make that gap permanent for the acquiring branch where the + full path would have filled it. The vectors are already in hand -- ``job.py`` + embeds every file the advisory read did not call unchanged. + """ + paths = [pf.path for pf, _sha in membership] + shas = [sha for _pf, sha in membership] + rows = conn.execute( + text( + "UPDATE files SET branches = (SELECT array_agg(DISTINCT e) FROM " + "unnest(files.branches || CAST(:branch_arr AS text[])) e) " + "WHERE repo_id = :repo_id " + "AND EXISTS (SELECT 1 FROM unnest(CAST(:paths AS text[]), CAST(:shas AS text[])) " + "AS t(p, s) WHERE t.p = files.path AND t.s = files.content_sha) " + "RETURNING id, path, content_sha" + ), + {"repo_id": repo_id, "branch_arr": [branch], "paths": paths, "shas": shas}, + ).all() + + if chunk_writer is None: + return + file_ids = {(row.path, row.content_sha): row.id for row in rows} + for pf, sha in membership: + file_id = file_ids.get((pf.path, sha)) + if file_id is None: + # Unreachable while the single-writer invariant holds: the row was in + # statement 3b's projection moments ago, inside this transaction. + logger.warning( + "membership-only row for %s vanished before its union; skipping its chunk write", + pf.path, + ) + continue + chunk_writer(conn, repo_id, file_id, pf) + + def _sweep_membership( conn: Connection, *, @@ -339,12 +634,43 @@ def _stamp_repo_branch( head_sha: str, baseline_commit: str | None, baseline_version: int | None, + seen_any: bool = True, ) -> None: """Compare-and-set the ``repo_branches`` stamp against the statement-2 baseline. Raises :class:`StaleIndexError` if the row no longer matches the baseline, which propagates out of ``index_repo``'s ``conn.begin()`` and rolls the whole ``(repo, branch)`` transaction back rather than regressing the index. + + **``seen_any=False`` holds the semantics version at ``baseline_version``.** + A run that parsed zero indexable files (the transient case + ``_sweep_membership``'s empty-seen-set guard exists for) has written nothing, + so it has not indexed anything at the CURRENT semantics version and must not + claim to have. Without this the following is silent and terminal: + + 1. ``INDEX_SEMANTICS_VERSION`` goes 4 -> 5; branch ``b`` is stored at + ``(sha1, 4)``, so the skip seam forces a re-index. + 2. That re-index parses zero files. Nothing is written -- but the stamp + advances to ``(sha2, 5)``. + 3. The next run sees ``baseline_version == 5 == current``, opens the + file-level delta gate, and finds every row carrying ``b`` unchanged on + ``(path, content_sha)`` -- so it skips them all. + 4. ``b`` serves v4-extracted rows under a v5 stamp, permanently. + + ``last_indexed_commit`` still advances to ``head_sha``: the commit IS what + this run looked at. Leaving the version behind is what makes the branch + mismatch (and therefore re-index) on its next run -- self-healing, in the + safe direction. The statement shape, the CAS predicate, and + :class:`StaleIndexError` are untouched. + + **Known, deliberate divergence:** ``index_repo``'s statement 1 writes the + DEPRECATED ``repos.index_semantics_version`` unconditionally on + ``is_default``, with no seen-set awareness. So a zero-parse default-branch + run leaves ``repos`` at the current version while ``repo_branches`` sits at + the old one. That is cosmetic -- no decision anywhere reads + ``repos.index_semantics_version`` (the three legacy columns are documented + deprecated in ``app/db/models.py``) -- and extending this fix to the legacy + stamp is scope this change deliberately does not take. """ result = conn.execute( update(RepoBranch) @@ -356,7 +682,7 @@ def _stamp_repo_branch( ) .values( last_indexed_commit=head_sha, - index_semantics_version=INDEX_SEMANTICS_VERSION, + index_semantics_version=(INDEX_SEMANTICS_VERSION if seen_any else baseline_version), last_indexed_at=func.now(), ) ) diff --git a/tests/integration/test_store.py b/tests/integration/test_store.py index 22cc75e..0cd4af1 100644 --- a/tests/integration/test_store.py +++ b/tests/integration/test_store.py @@ -136,16 +136,41 @@ def test_first_run_populates_and_stamps_commit(conn: Connection) -> None: @pytest.mark.integration def test_rerun_is_idempotent(conn: Connection) -> None: + """The first run stamps INDEX_SEMANTICS_VERSION, so the second run (identical + content, identical head_sha) hits the file-level delta gate: both files are + classified unchanged, so `symbols` in the returned IndexCounts is 0 (files + seen, not files re-inserted) -- not the pre-delta re-insert count of 2. See + test_rerun_is_idempotent_preserves_row_identity below for the stronger + row-identity property this test was originally reaching for. + """ _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(MAIN, UTIL)) counts = _index_default( conn, name="acme/widgets", head_sha="sha_first", items=_items(MAIN, UTIL) ) - assert counts == IndexCounts(files=2, symbols=2, swept=0, edges=0) + assert counts == IndexCounts(files=2, symbols=0, swept=0, edges=0) assert _count(conn, "repos") == 1 assert _count(conn, "files") == 2 assert _count(conn, "symbols") == 2 +@pytest.mark.integration +def test_rerun_is_idempotent_preserves_row_identity(conn: Connection) -> None: + """The stronger property test_rerun_is_idempotent was originally reaching + for: an unchanged re-run does not delete-and-reinsert ANYTHING -- files.id + and symbols.id survive byte-identical across the two runs (a delete-reinsert + would renumber the serials).""" + _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(MAIN, UTIL)) + files_before = dict(conn.execute(text("SELECT path, id FROM files ORDER BY path")).all()) + symbols_before = sorted(conn.execute(text("SELECT id FROM symbols")).scalars().all()) + conn.rollback() + + _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(MAIN, UTIL)) + files_after = dict(conn.execute(text("SELECT path, id FROM files ORDER BY path")).all()) + symbols_after = sorted(conn.execute(text("SELECT id FROM symbols")).scalars().all()) + assert files_after == files_before + assert symbols_after == symbols_before + + @pytest.mark.integration def test_mark_and_sweep_removes_deleted_file(conn: Connection) -> None: # util.py has a real call site so the writer produces a reference_edges row @@ -168,15 +193,23 @@ def test_mark_and_sweep_removes_deleted_file(conn: Connection) -> None: # (production hands a fresh engine.connect() per repo). conn.rollback() - # Re-run without util.py and with a new head SHA -> util.py is swept. + # Re-run without util.py and with a new head SHA -> util.py is swept. MAIN's + # content is identical across both runs, and the first run already stamped + # INDEX_SEMANTICS_VERSION, so the delta gate classifies MAIN unchanged: no + # symbol rewrite (symbols=0, not a re-insert count). counts = _index_default(conn, name="acme/widgets", head_sha="sha_second", items=_items(MAIN)) - assert counts == IndexCounts(files=1, symbols=1, swept=1, edges=0) + assert counts == IndexCounts(files=1, symbols=0, swept=1, edges=0) assert _count(conn, "files", "path = 'util.py'") == 0 assert _count(conn, "symbols", f"file_id = {removed_file_id}") == 0 # cascade assert _count(conn, "reference_edges", f"file_id = {removed_file_id}") == 0 # cascade assert _count(conn, "files") == 1 - assert _count(conn, "files", "commit = 'sha_second'") == 1 + # MAIN's row was classified unchanged (skipped), so its `commit` column stays + # at the FIRST run's SHA -- files.commit goes staler under delta indexing by + # design (no production read path resolves `commit:` from it; see + # indexer/store.py's index_repo docstring). + assert _count(conn, "files", "commit = 'sha_first'") == 1 + assert _count(conn, "files", "commit = 'sha_second'") == 0 @pytest.mark.integration @@ -196,14 +229,15 @@ def test_index_repo_records_the_sweep_phase(conn: Connection) -> None: token = install_timer(timer) try: # Re-run without util.py at a new SHA: the same scenario as - # test_mark_and_sweep_removes_deleted_file, so the counts are unchanged. + # test_mark_and_sweep_removes_deleted_file, so the counts are unchanged + # (MAIN classified unchanged under the delta gate -> symbols=0). counts = _index_default( conn, name="acme/widgets", head_sha="sha_second", items=_items(MAIN) ) finally: reset_timer(token) - assert counts == IndexCounts(files=1, symbols=1, swept=1, edges=0) + assert counts == IndexCounts(files=1, symbols=0, swept=1, edges=0) assert timer.total("sweep") > 0.0 # index_repo measures the sweep and nothing else -- every other phase is # job.py's to record. @@ -224,8 +258,10 @@ def test_sweep_is_repo_scoped(conn: Connection) -> None: conn.rollback() # Re-index A without util.py at a new SHA -> A's util.py swept, B untouched. + # MAIN unchanged under the delta gate (A's first run already stamped + # INDEX_SEMANTICS_VERSION) -> symbols=0, same as the tests above. counts = _index_default(conn, name="acme/a", head_sha="a_second", items=_items(MAIN)) - assert counts == IndexCounts(files=1, symbols=1, swept=1, edges=0) + assert counts == IndexCounts(files=1, symbols=0, swept=1, edges=0) assert _count(conn, "files", "repo_id = (SELECT id FROM repos WHERE name = 'acme/b')") == ( b_files_before @@ -267,6 +303,17 @@ def test_reindex_replaces_stale_edges_for_the_same_file(conn: Connection) -> Non deleted and reinserted, not accumulated -- proven by driving the two runs' ``ex.edges`` directly rather than depending on the real extractor to disagree with itself on unchanged content. + + This is precisely the case the file-level delta gate (#104) would otherwise + skip: identical content_sha, different extraction output. In production that + combination can only arise from an extractor change, which mandates an + INDEX_SEMANTICS_VERSION bump and therefore closes the gate on its own -- so + forcing it closed here (the same ``UPDATE repo_branches SET + index_semantics_version = NULL`` idiom as + test_legacy_null_semantics_version_is_rewritten) is faithful to production, + not a workaround. This is the "unconditional-delete guard" and it must keep + its ORIGINAL assertion, not a relaxed one -- see the plan for issue #104, + §2.6a. """ symbol = ExtractedSymbol("f", "function", 1, 3) content = "def f():\n target()\n return 1\n" @@ -279,7 +326,8 @@ def test_reindex_replaces_stale_edges_for_the_same_file(conn: Connection) -> Non ) file_id = conn.execute(text("SELECT id FROM files WHERE path = 'main.py'")).scalar_one() assert _count(conn, "reference_edges", f"file_id = {file_id}") == 1 - conn.rollback() + conn.execute(text("UPDATE repo_branches SET index_semantics_version = NULL")) + conn.commit() second_edge = ExtractedEdge(kind="call", target="new_target", line=2, enclosing=symbol) _index_default( @@ -303,6 +351,15 @@ def test_reindex_replaces_stale_edges_for_the_same_file(conn: Connection) -> Non @pytest.mark.integration def test_reindex_with_identical_items_does_not_duplicate_edges(conn: Connection) -> None: + """The delete-before-insert idempotency of the edge writer, proven by forcing + the delta gate CLOSED (the same idiom as the two guard tests above) so the + second run genuinely re-executes the write path rather than classifying + main.py unchanged and skipping it -- which would make the `== 1` assertion + pass vacuously (nothing touched at all) rather than proving delete-then- + insert doesn't duplicate. Without this the first run's first-index-ever + baseline (version None) makes run 2's baseline INDEX_SEMANTICS_VERSION, and + the delta gate would otherwise open and skip main.py entirely. + """ symbol = ExtractedSymbol("f", "function", 1, 3) item = ( "main.py", @@ -311,7 +368,8 @@ def test_reindex_with_identical_items_does_not_duplicate_edges(conn: Connection) [ExtractedEdge(kind="call", target="helper", line=2, enclosing=symbol)], ) _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(item)) - conn.rollback() + conn.execute(text("UPDATE repo_branches SET index_semantics_version = NULL")) + conn.commit() _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(item)) file_id = conn.execute(text("SELECT id FROM files WHERE path = 'main.py'")).scalar_one() @@ -326,6 +384,12 @@ def test_reindex_to_zero_edges_sheds_all_rows(conn: Connection) -> None: upsert resolves to the SAME ``file_id`` -- isolating the write-side guard (``ex.edges`` empty must still run the delete) from the unrelated delete-and-reinsert-under-a-new-file-id path already covered above. + + Same rationale as test_reindex_replaces_stale_edges_for_the_same_file: this + is identical content_sha with divergent extraction output, which in + production can only happen via an extractor change (mandating a semantics + version bump). The gate is forced closed here rather than the expected + values relaxed -- see the plan for issue #104, §2.6a. """ symbol = ExtractedSymbol("f", "function", 1, 3) content = "def f():\n helper()\n return 1\n" @@ -338,7 +402,8 @@ def test_reindex_to_zero_edges_sheds_all_rows(conn: Connection) -> None: ) file_id = conn.execute(text("SELECT id FROM files WHERE path = 'main.py'")).scalar_one() assert _count(conn, "reference_edges", f"file_id = {file_id}") == 1 - conn.rollback() + conn.execute(text("UPDATE repo_branches SET index_semantics_version = NULL")) + conn.commit() # Same content (same file_id) but this run's extraction yields zero edges. counts = _index_default( @@ -598,7 +663,8 @@ def test_per_branch_cas_resume_is_independent_per_branch(conn: Connection) -> No conn.rollback() # Re-indexing 'a' again must succeed against its own baseline, unaffected by - # 'b' having indexed in between. + # 'b' having indexed in between. MAIN unchanged under the delta gate ('a's + # first run already stamped INDEX_SEMANTICS_VERSION) -> symbols=0. counts = index_repo( conn, name="acme/widgets", @@ -607,7 +673,7 @@ def test_per_branch_cas_resume_is_independent_per_branch(conn: Connection) -> No head_sha="sha_a2", items=_items(MAIN), ) - assert counts == IndexCounts(files=1, symbols=1, swept=0, edges=0) + assert counts == IndexCounts(files=1, symbols=0, swept=0, edges=0) stamps = dict(conn.execute(text("SELECT branch, last_indexed_commit FROM repo_branches")).all()) assert stamps == {"a": "sha_a2", "b": "sha_b"} diff --git a/tests/integration/test_store_chunk_writer.py b/tests/integration/test_store_chunk_writer.py index 0dd67f5..30815b4 100644 --- a/tests/integration/test_store_chunk_writer.py +++ b/tests/integration/test_store_chunk_writer.py @@ -12,6 +12,17 @@ chunk_writer ride the same conn.begin() as the rest of that file's row, and cascade-delete when the file is swept (FK ON DELETE CASCADE), exactly like symbols. + +**Every test in this module is Lakebase-deferred for issue #104**: this is the +one module whose fixture builds the ``chunks`` table and needs +``lakebase_vector``, which no local Postgres image provides (see +``tests/integration/test_store_delta.py``'s module docstring for why the +core delta suite deliberately lives elsewhere instead). The delta-specific +additions here (the unchanged-file / membership-only chunk cases, and +``test_reindex_is_idempotent_for_chunks``'s reviewed expectations) are +reasoned through against ``indexer/store.py``'s ``_union_membership`` and +the per-file loop, never verified by a local run -- flagged explicitly in the +PR body, per the plan for issue #104, §2.6a / §3.2. """ from __future__ import annotations @@ -136,6 +147,19 @@ def test_chunk_writer_writes_inside_the_transaction(conn: Connection) -> None: @pytest.mark.integration def test_reindex_is_idempotent_for_chunks(conn: Connection) -> None: + """Reviewed against the plan for issue #104, §2.6a's treatment table (Lakebase- + deferred, so this review is reasoned through rather than locally verified -- + see the module docstring). Unlike the two GUARD tests in + ``tests/integration/test_store.py`` (which deliberately hold content_sha + identical while varying EXTRACTION output, and must keep their gate forced + closed), this run is identical in every respect -- content, head_sha, AND + extraction. The first run stamps INDEX_SEMANTICS_VERSION, so the second + classifies both files unchanged: ``symbols`` drops from 2 (a delete-reinsert + count) to 0 (nothing rewritten), a plain (a)-style value update. The chunk + row COUNT is unaffected either way -- 2 rows survive whether by an idempotent + delete-reinsert (pre-#104) or by never being touched at all (unchanged, under + #104) -- so that assertion needed no change, only its reasoning. + """ items = _items(MAIN, UTIL) index_repo( conn, @@ -156,8 +180,8 @@ def test_reindex_is_idempotent_for_chunks(conn: Connection) -> None: items=items, chunk_writer=_stub_chunk_writer, ) - assert counts == IndexCounts(files=2, symbols=2, swept=0, edges=0) - assert _count(conn, "chunks") == 2 # delete-and-reinsert, not duplicated + assert counts == IndexCounts(files=2, symbols=0, swept=0, edges=0) + assert _count(conn, "chunks") == 2 # untouched, not delete-and-reinserted @pytest.mark.integration @@ -191,3 +215,100 @@ def test_chunks_cascade_delete_when_file_is_swept(conn: Connection) -> None: assert _count(conn, "files", "path = 'util.py'") == 0 assert _count(conn, "chunks", f"file_id = {util_file_id}") == 0 # cascade assert _count(conn, "chunks") == 1 + + +# --- File-level delta indexing (#104): unchanged / membership-only chunks --- +# Lakebase-deferred -- see the module docstring. Reasoned through against +# indexer/store.py's per-file loop and _union_membership, never locally run. + + +@pytest.mark.integration +def test_unchanged_file_never_calls_chunk_writer_and_preserves_chunk_ids( + conn: Connection, +) -> None: + """An unchanged file's chunks.id values are IDENTICAL before/after, and + chunk_writer is never called for it at all -- proven with a call-tracking + wrapper, not just by the row count staying flat (which an idempotent + delete-reinsert of identical content would also produce, as + test_reindex_is_idempotent_for_chunks above shows).""" + calls: list[str] = [] + + def _tracking_chunk_writer( + conn: Connection, repo_id: int, file_id: int, pf: ParsedFile + ) -> None: + calls.append(pf.path) + _stub_chunk_writer(conn, repo_id, file_id, pf) + + index_repo( + conn, + name="acme/widgets", + branch="main", + is_default=True, + head_sha="sha_first", + items=_items(MAIN, UTIL), + chunk_writer=_tracking_chunk_writer, + ) + chunk_ids_before = sorted(conn.execute(text("SELECT id FROM chunks")).scalars().all()) + conn.rollback() + calls.clear() + + index_repo( + conn, + name="acme/widgets", + branch="main", + is_default=True, + head_sha="sha_second", + items=_items(MAIN, UTIL), + chunk_writer=_tracking_chunk_writer, + ) + assert calls == [] + chunk_ids_after = sorted(conn.execute(text("SELECT id FROM chunks")).scalars().all()) + assert chunk_ids_after == chunk_ids_before + + +@pytest.mark.integration +def test_membership_only_file_backfills_previously_missing_chunks(conn: Connection) -> None: + """§2.4 note 3: a membership-acquired file whose row had ZERO chunk rows + (branch 'a' wrote it with chunk_writer=None, i.e. semantic-off) ends branch + 'b's acquiring run WITH chunk rows. Membership-only writes chunks even + though it writes no symbols/edges -- the vectors are already in hand + (job.py embeds every file the advisory read did not call unchanged), so + skipping the write would make a semantic-off-then-on transition's gap + permanent for the acquiring branch instead of backfilling it.""" + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a1", + items=_items(MAIN), # no chunk_writer -> main.py has zero chunk rows + ) + conn.rollback() + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b1", + items=_items(UTIL), + chunk_writer=_stub_chunk_writer, + ) + main_file_id = conn.execute(text("SELECT id FROM files WHERE path = 'main.py'")).scalar_one() + assert _count(conn, "chunks", f"file_id = {main_file_id}") == 0 + conn.rollback() + + # branch 'b' is now at the current semantics version (its own baseline), + # so acquiring main.py (identical content, still stored under 'a') takes + # the membership-only path. + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b2", + items=_items(MAIN, UTIL), + chunk_writer=_stub_chunk_writer, + ) + assert _count(conn, "chunks", f"file_id = {main_file_id}") == 1 + branches = conn.execute(text("SELECT branches FROM files WHERE path = 'main.py'")).scalar_one() + assert sorted(branches) == ["a", "b"] diff --git a/tests/integration/test_store_delta.py b/tests/integration/test_store_delta.py new file mode 100644 index 0000000..08e528a --- /dev/null +++ b/tests/integration/test_store_delta.py @@ -0,0 +1,658 @@ +"""Integration tests for the file-level delta path (issue #104) against real Postgres. + +Row-identity proof -- that skipping a file really does leave its stored serials +untouched -- is this module's job; ``tests/unit/test_store_delta.py`` pins the +exact *statement inventory* each classification produces against a fake +connection. Clones ``tests/integration/test_store.py``'s throwaway-schema fixture +idiom (own copy, per ``tests/integration/AGENTS.md``'s no-conftest convention): +a clean schema per run, the durable-core DDL via ``Base.metadata.create_all``, +and a per-connection ``search_path`` that propagates into ``index_repo``'s DML. + +**Deliberately builds no ``chunks`` table and creates no ``lakebase_*`` +extension.** That is the entire reason this module runs locally at all -- +``test_store_chunk_writer.py`` errors at *module fixture* setup on +``CREATE EXTENSION IF NOT EXISTS lakebase_vector CASCADE``, which no local +Postgres image provides. Chunk-touching delta cases live there instead +(Lakebase-deferred; see that module's docstring). +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Iterator +from typing import Any + +import pytest +from sqlalchemy import Connection, text + +from app.db.client import create_db_engine +from app.db.models import INDEX_SEMANTICS_VERSION, Base +from indexer.languages import ExtractedSymbol, FileExtraction, IndexCounts, ParsedFile +from indexer.store import index_repo + +SCHEMA = "test_store_delta" + + +@pytest.fixture +def conn() -> Iterator[Connection]: + engine = create_db_engine() + connection = engine.connect() + try: + connection.execute(text(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE")) + connection.execute(text(f"CREATE SCHEMA {SCHEMA}")) + connection.execute(text(f"SET search_path TO {SCHEMA}, public")) + connection.commit() + + Base.metadata.create_all(bind=connection) + connection.commit() + + yield connection + finally: + connection.rollback() + connection.execute(text(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE")) + connection.commit() + connection.close() + engine.dispose() + + +def _pf(path: str, content: str) -> ParsedFile: + return ParsedFile(path=path, lang="python", size=len(content.encode()), content=content) + + +def _items( + *specs: tuple[str, str, list[ExtractedSymbol]], +) -> list[tuple[ParsedFile, FileExtraction]]: + return [ + (_pf(path, content), FileExtraction(symbols=syms, edges=[])) + for path, content, syms in specs + ] + + +def _fn(name: str, n: int) -> tuple[str, str, list[ExtractedSymbol]]: + """One deterministic (path, content, symbols) triple, distinguishable by ``n``.""" + content = f"def f{n}():\n return {n}\n" + return (f"{name}{n}.py", content, [ExtractedSymbol(f"f{n}", "function", 1, 2)]) + + +MAIN = ("main.py", "def f():\n return 1\n", [ExtractedSymbol("f", "function", 1, 2)]) +UTIL = ("util.py", "def g():\n return 2\n", [ExtractedSymbol("g", "function", 1, 2)]) + + +def _index_default( + conn: Connection, *, name: str, head_sha: str, items: list[tuple[ParsedFile, FileExtraction]] +) -> IndexCounts: + return index_repo( + conn, name=name, branch="main", is_default=True, head_sha=head_sha, items=items + ) + + +def _count(conn: Connection, table: str, where: str = "") -> int: + sql = f"SELECT count(*) FROM {table}" + if where: + sql += f" WHERE {where}" + return int(conn.execute(text(sql)).scalar_one()) + + +def _symbol_ids(conn: Connection, path: str) -> list[int]: + return sorted( + conn.execute( + text("SELECT s.id FROM symbols s JOIN files f ON f.id = s.file_id WHERE f.path = :p"), + {"p": path}, + ) + .scalars() + .all() + ) + + +def _delta_lines(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + r.getMessage() + for r in caplog.records + if r.name == "indexer.store" and "delta write set" in r.getMessage() + ] + + +# --- Test 15: row-identity proof, the "unchanged" class writes NOTHING ------ + + +@pytest.mark.integration +def test_unchanged_rerun_preserves_every_row_identity( + conn: Connection, caplog: pytest.LogCaptureFixture +) -> None: + """Re-index identical content at a NEW head_sha: files.id and symbols.id are + IDENTICAL before/after -- a delete-reinsert would renumber the serials, so + identical ids are the precise proof that zero rows were written. + ``counts.symbols == 0`` while ``counts.files == N``, and the + ``delta write set 0/N`` line is emitted. + """ + items = _items(MAIN, UTIL) + _index_default(conn, name="acme/widgets", head_sha="sha_first", items=items) + files_before = dict(conn.execute(text("SELECT path, id FROM files ORDER BY path")).all()) + symbols_before = sorted(conn.execute(text("SELECT id FROM symbols")).scalars().all()) + conn.rollback() + + with caplog.at_level(logging.INFO, logger="indexer.store"): + counts = _index_default( + conn, name="acme/widgets", head_sha="sha_second", items=_items(MAIN, UTIL) + ) + assert counts == IndexCounts(files=2, symbols=0, swept=0, edges=0) + + files_after = dict(conn.execute(text("SELECT path, id FROM files ORDER BY path")).all()) + symbols_after = sorted(conn.execute(text("SELECT id FROM symbols")).scalars().all()) + assert files_after == files_before + assert symbols_after == symbols_before + + lines = _delta_lines(caplog) + assert len(lines) == 1 + assert "delta write set 0/2 files (unchanged=2 membership=0, semantics gate open)" in lines[0] + + +# --- Test 16: small delta (1 changed, 1 added, 1 deleted) -- issue AC2 ------- + + +@pytest.mark.integration +def test_small_delta_writes_only_the_changed_added_and_sweeps_the_deleted(conn: Connection) -> None: + """1 changed, 1 unchanged, 1 added, 1 deleted -- the changed file's symbol ids + change, the unchanged file's do not, the added file is present, and the + deleted file is swept. Issue #104's acceptance criterion 2. + """ + unchanged_symbol = ExtractedSymbol("f", "function", 1, 2) + changed_v1 = ( + "changed.py", + "def c():\n return 1\n", + [ExtractedSymbol("c", "function", 1, 2)], + ) + to_delete = ("gone.py", "def d():\n return 1\n", [ExtractedSymbol("d", "function", 1, 2)]) + _index_default( + conn, + name="acme/widgets", + head_sha="sha_first", + items=_items( + ("unchanged.py", "def f():\n return 1\n", [unchanged_symbol]), changed_v1, to_delete + ), + ) + unchanged_ids_before = _symbol_ids(conn, "unchanged.py") + changed_ids_before = _symbol_ids(conn, "changed.py") + assert unchanged_ids_before and changed_ids_before + conn.rollback() + + changed_v2 = ( + "changed.py", + "def c():\n return 2\n", + [ExtractedSymbol("c", "function", 1, 2)], + ) + added = ("added.py", "def a():\n return 1\n", [ExtractedSymbol("a", "function", 1, 2)]) + counts = _index_default( + conn, + name="acme/widgets", + head_sha="sha_second", + items=_items( + ("unchanged.py", "def f():\n return 1\n", [unchanged_symbol]), changed_v2, added + ), + ) + # swept=2: gone.py (removed outright) AND changed.py's OLD content_sha row + # (the changed-file upsert mints a NEW row under the new content_sha, since + # uq_files_repo_path_sha is keyed on content -- the stale old-sha row is + # exactly what the sweep exists to reap, unrelated to delta). + assert counts == IndexCounts(files=3, symbols=2, swept=2, edges=0) + + assert _symbol_ids(conn, "unchanged.py") == unchanged_ids_before + assert _symbol_ids(conn, "changed.py") != changed_ids_before + assert _count(conn, "files", "path = 'added.py'") == 1 + assert _count(conn, "files", "path = 'gone.py'") == 0 + + +# --- Test 17: multi-branch dedup takes the membership-only path -- AC3 ------ + + +@pytest.mark.integration +def test_membership_only_dedup_across_branches_preserves_symbol_ids( + conn: Connection, caplog: pytest.LogCaptureFixture +) -> None: + """Branch 'b' acquires content already stored under branch 'a': ``branches`` + becomes ``['a', 'b']`` (sorted, matching the array_agg(DISTINCT ...) idiom), + symbol ids are UNCHANGED (no rewrite), and the delta write set line reports + ``membership=1``. Issue #104's acceptance criterion 3. + """ + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a1", + items=_items(MAIN), + ) + conn.rollback() + # branch 'b' must ALREADY be at the current semantics version (its own + # baseline) for the delta gate to be open when it next acquires MAIN -- + # a brand-new branch's first run is always full-path. + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b1", + items=_items(UTIL), + ) + main_ids_before = _symbol_ids(conn, "main.py") + assert main_ids_before + conn.rollback() + + with caplog.at_level(logging.INFO, logger="indexer.store"): + counts = index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b2", + items=_items(MAIN, UTIL), + ) + assert counts == IndexCounts(files=2, symbols=0, swept=0, edges=0) + + assert _symbol_ids(conn, "main.py") == main_ids_before + branches = conn.execute(text("SELECT branches FROM files WHERE path = 'main.py'")).scalar_one() + assert sorted(branches) == ["a", "b"] + + lines = _delta_lines(caplog) + assert len(lines) == 1 + assert "membership=1" in lines[0] + + +@pytest.mark.integration +def test_membership_only_row_gets_no_symbol_or_edge_statement(conn: Connection) -> None: + """Companion to the row-identity proof above, phrased as a statement-absence + check rather than an id-equality check: the acquiring branch's run inserts + NO symbols row for the acquired file at all (there was never a duplicate to + delete-and-reinsert).""" + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a1", + items=_items(MAIN), + ) + conn.rollback() + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b1", + items=_items(UTIL), + ) + symbols_before = _count(conn, "symbols") + conn.rollback() + + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b2", + items=_items(MAIN, UTIL), + ) + assert _count(conn, "symbols") == symbols_before + + +# --- Test 18: a semantics-version mismatch forces the full path -- AC4 ------ + + +@pytest.mark.integration +def test_stale_semantics_version_forces_the_full_path_for_every_file(conn: Connection) -> None: + _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(MAIN, UTIL)) + main_ids_before = _symbol_ids(conn, "main.py") + util_ids_before = _symbol_ids(conn, "util.py") + conn.execute( + text("UPDATE repo_branches SET index_semantics_version = :v"), + {"v": INDEX_SEMANTICS_VERSION - 1}, + ) + conn.commit() + + counts = _index_default( + conn, name="acme/widgets", head_sha="sha_second", items=_items(MAIN, UTIL) + ) + assert counts == IndexCounts(files=2, symbols=2, swept=0, edges=0) + assert _symbol_ids(conn, "main.py") != main_ids_before + assert _symbol_ids(conn, "util.py") != util_ids_before + + +# --- Test 19: the provenance gate -- a stale SIBLING branch closes it ------- + + +@pytest.mark.integration +def test_provenance_gate_forces_the_full_path_when_a_sibling_branch_is_stale( + conn: Connection, +) -> None: + """The counter-example the provenance gate (statement 4) exists to close: + 'stale_sibling' wrote main.py at the current version, then regressed to an + older one (simulating a failed re-index). 'acquirer' is itself at the + current version and tries to acquire main.py membership-only -- but every + ``repo_branches`` row for this repo must be current for that path to be + taken, and stale_sibling's is not, so 'acquirer' gets the FULL path + instead (proven by main.py's symbol ids changing under 'acquirer'). + """ + index_repo( + conn, + name="acme/widgets", + branch="stale_sibling", + is_default=True, + head_sha="sha_s1", + items=_items(MAIN), + ) + main_ids_before = _symbol_ids(conn, "main.py") + conn.execute( + text( + "UPDATE repo_branches SET index_semantics_version = :v WHERE branch = 'stale_sibling'" + ), + {"v": INDEX_SEMANTICS_VERSION - 1}, + ) + conn.commit() + + index_repo( + conn, + name="acme/widgets", + branch="acquirer", + is_default=False, + head_sha="sha_a1", + items=_items(UTIL), + ) + conn.rollback() + + index_repo( + conn, + name="acme/widgets", + branch="acquirer", + is_default=False, + head_sha="sha_a2", + items=_items(MAIN, UTIL), + ) + + assert _symbol_ids(conn, "main.py") != main_ids_before + branches = conn.execute(text("SELECT branches FROM files WHERE path = 'main.py'")).scalar_one() + assert sorted(branches) == ["acquirer", "stale_sibling"] + + +# --- Tests 20a/20b: the membership invariant, both directions -------------- + + +@pytest.mark.integration +def test_no_files_row_ever_has_an_empty_branches_array(conn: Connection) -> None: + """Forward direction: both sweep sites (the per-branch sweep and the + array-remove path) delete a row once its branches array is emptied, never + leaving a zombie row behind. Exercised across additions, a shared-then- + removed-from-one-branch file, and a fully-removed file. + """ + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a1", + items=_items(MAIN, UTIL), + ) + conn.rollback() + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b1", + items=_items(UTIL), + ) + conn.rollback() + # Empty-seen-set guard fires (see store.py's _sweep_membership): a NO-OP, + # not a drop. Included to prove the guard doesn't itself leave or create an + # empty-branches row. + index_repo(conn, name="acme/widgets", branch="a", is_default=True, head_sha="sha_a2", items=[]) + conn.rollback() + # 'a' now genuinely drops both files (UTIL is the only file it still + # parses): util.py loses only 'a' (row survives, still shared with 'b'); + # main.py loses its only remaining branch and is deleted outright. + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a3", + items=_items(UTIL), + ) + conn.rollback() + # Empty-seen-set guard again, now that 'a' carries nothing -- still a no-op. + index_repo(conn, name="acme/widgets", branch="a", is_default=True, head_sha="sha_a4", items=[]) + + assert _count(conn, "files", "cardinality(branches) = 0") == 0 + + +@pytest.mark.integration +def test_every_branch_on_a_files_row_has_a_repo_branches_row(conn: Connection) -> None: + """Converse direction: index_repo writes statement 2 (the repo_branches + upsert) before any file row for that branch, so no branches array element + can ever dangle without a matching repo_branches row. The provenance gate + (statement 4) depends on this holding.""" + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a1", + items=_items(MAIN), + ) + conn.rollback() + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b1", + items=_items(MAIN, UTIL), + ) + + dangling = conn.execute( + text( + "SELECT count(*) FROM files f, unnest(f.branches) AS b " + "WHERE NOT EXISTS (" + " SELECT 1 FROM repo_branches rb WHERE rb.repo_id = f.repo_id AND rb.branch = b" + ")" + ) + ).scalar_one() + assert dangling == 0 + + +# --- Test 21: empty-seen-set guard and per-branch CAS still hold with delta - + + +@pytest.mark.integration +def test_empty_seen_set_guard_holds_with_the_delta_gate_open(conn: Connection) -> None: + """The empty-seen-set guard (skip the sweep, WARN, return 0) is delta-blind by + construction -- it fires on ``seen_paths`` being empty, before any + classification runs. Pinned again here because it is exactly the run shape + the delta gate makes common (a branch whose HEAD moved but touched nothing + indexable).""" + _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(MAIN)) + conn.rollback() + # Second run: delta gate is open (first run stamped INDEX_SEMANTICS_VERSION), + # but this run parses zero files. + counts = _index_default(conn, name="acme/widgets", head_sha="sha_second", items=[]) + assert counts == IndexCounts(files=0, symbols=0, swept=0, edges=0) + assert _count(conn, "files", "path = 'main.py'") == 1 + branches = conn.execute(text("SELECT branches FROM files WHERE path = 'main.py'")).scalar_one() + assert branches == ["main"] + + +@pytest.mark.integration +def test_cas_still_rejects_a_stale_baseline_with_delta_on(conn: Connection) -> None: + """The CAS predicate is statement 5, downstream of every delta statement -- + proven still load-bearing by forcing a conflict on a branch whose delta gate + is open (second run onward).""" + from indexer.store import StaleIndexError, _stamp_repo_branch + + _index_default(conn, name="acme/widgets", head_sha="sha_a", items=_items(MAIN, UTIL)) + conn.rollback() + _index_default(conn, name="acme/widgets", head_sha="sha_b", items=_items(MAIN, UTIL)) + files_before = _count(conn, "files") + repo_id = int( + conn.execute(text("SELECT id FROM repos WHERE name = 'acme/widgets'")).scalar_one() + ) + conn.rollback() + + with pytest.raises(StaleIndexError, match="wrong_sha"), conn.begin(): + conn.execute(text("DELETE FROM files WHERE path = 'util.py'")) + _stamp_repo_branch( + conn, + name="acme/widgets", + branch="main", + repo_id=repo_id, + head_sha="sha_c", + baseline_commit="wrong_sha", + baseline_version=INDEX_SEMANTICS_VERSION, + ) + assert _count(conn, "files") == files_before + + +# --- Test 22 (BLOCKER): a zero-parse run must NOT advance the semantics stamp + + +@pytest.mark.integration +def test_zero_parse_run_does_not_advance_the_semantics_stamp(conn: Connection) -> None: + """The base-case fix (§2.3): index a branch non-empty, force its stored + version DOWN to simulate a pre-transition stamp, then re-index with + ``items=[]`` at a NEW head SHA. The stored version must stay at the forced + value (NOT advance to INDEX_SEMANTICS_VERSION) -- proving a zero-parse run + cannot manufacture a spurious "current version" base case for the delta + induction. The commit still advances (the run DID look at that SHA); only + the version is held back. The FOLLOWING non-empty run must then full-path + every file, since the branch is still stamped stale. + """ + _index_default(conn, name="acme/widgets", head_sha="sha_first", items=_items(MAIN, UTIL)) + main_ids_before = _symbol_ids(conn, "main.py") + conn.execute(text("UPDATE repo_branches SET index_semantics_version = 3")) + conn.commit() + + counts = _index_default(conn, name="acme/widgets", head_sha="sha_zero", items=[]) + assert counts == IndexCounts(files=0, symbols=0, swept=0, edges=0) + stamp = conn.execute( + text( + "SELECT rb.last_indexed_commit, rb.index_semantics_version FROM repo_branches rb " + "JOIN repos r ON r.id = rb.repo_id WHERE r.name = 'acme/widgets'" + ) + ).one() + # last_indexed_commit DID advance (the run looked at sha_zero); the version + # did NOT (nothing was indexed at it). + assert stamp == ("sha_zero", 3) + conn.rollback() + + # The next non-empty run sees baseline_version=3 != INDEX_SEMANTICS_VERSION, + # so the gate is closed and every file takes the full path -- symbol ids + # change even though the content is byte-identical to sha_first's. + _index_default(conn, name="acme/widgets", head_sha="sha_second", items=_items(MAIN, UTIL)) + assert _symbol_ids(conn, "main.py") != main_ids_before + stamp_after = conn.execute( + text( + "SELECT rb.index_semantics_version FROM repo_branches rb " + "JOIN repos r ON r.id = rb.repo_id WHERE r.name = 'acme/widgets'" + ) + ).scalar_one() + assert stamp_after == INDEX_SEMANTICS_VERSION + + +# --- Test 23: the pre-read is index-served, not a Seq Scan ------------------ + + +def _explain( + conn: Connection, sql: str, params: dict[str, Any], *, analyze: bool = False +) -> dict[str, Any]: + """``EXPLAIN (FORMAT JSON)`` for ``sql`` with the seq-scan escape hatch disabled, + scoped to a SAVEPOINT so the ``enable_seqscan`` GUC change never leaks past this + call (clones ``tests/integration/test_query_compiler.py``'s ``_explain_plan`` + idiom). ``VERBOSE`` is load-bearing, not decoration: Postgres only emits each + node's ``Output`` column list under ``VERBOSE``, and the projected-columns + assertions below depend on that field actually being present rather than + silently absent (which would make them vacuously true).""" + mode = "VERBOSE, ANALYZE, FORMAT JSON" if analyze else "VERBOSE, FORMAT JSON" + savepoint = conn.begin_nested() + try: + conn.execute(text("SET LOCAL enable_seqscan = off")) + raw = conn.execute(text(f"EXPLAIN ({mode}) {sql}"), params).scalar_one() + finally: + savepoint.rollback() + plan_list = json.loads(raw) if isinstance(raw, str) else raw + plan: dict[str, Any] = plan_list[0]["Plan"] + return plan + + +def _projects_column(output: list[str] | None, column: str) -> bool: + """Does EXPLAIN VERBOSE's ``Output`` list include ``column``, exactly or + table-qualified (``files.content``)? A plain substring check would false- + positive on ``content_sha`` containing ``content`` -- this checks whole + output entries instead.""" + return any(entry == column or entry.endswith(f".{column}") for entry in output or []) + + +def _plan_nodes(plan: dict[str, Any]) -> Iterator[dict[str, Any]]: + yield plan + for child in plan.get("Plans", []) or []: + yield from _plan_nodes(child) + + +@pytest.mark.integration +def test_the_delta_pre_read_is_index_served(conn: Connection) -> None: + """Statement 3a (``branches @> ...``) plans an index scan on + ``ix_files_branches_gin``; statement 3b (unqualified, repo-scoped) plans an + Index Only Scan on ``uq_files_repo_path_sha`` with zero heap fetches -- + proving ``read_repo_content_shas`` never resorts to a Seq Scan and never + pulls ``content`` off the heap. Seeded with 200 rows and ``VACUUM ANALYZE``d + (an Index Only Scan additionally needs a set visibility map, which + freshly-inserted, never-vacuumed rows do not have) so the plan is not a + tiny-corpus degenerate choice -- see + ``test_query_compiler.py::_explain_plan``'s docstring for that failure mode. + + ``VACUUM`` cannot run inside a transaction block, and SQLAlchemy 2.0 + autobegins one on every ``execute`` (a preceding ``conn.commit()`` does not + help -- the next ``execute`` autobegins again), so it runs on a SEPARATE + connection with ``execution_options(isolation_level="AUTOCOMMIT")``, never + on this module's transaction-bound ``conn`` fixture. + """ + items = _items(*[_fn("f", i) for i in range(200)]) + _index_default(conn, name="acme/widgets", head_sha="sha_first", items=items) + repo_id = int( + conn.execute(text("SELECT id FROM repos WHERE name = 'acme/widgets'")).scalar_one() + ) + conn.commit() + + engine = conn.engine + with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as vac_conn: + vac_conn.execute(text(f"SET search_path TO {SCHEMA}, public")) + vac_conn.execute(text("VACUUM ANALYZE files")) + + plan_a = _explain( + conn, + "SELECT path, content_sha FROM files " + "WHERE repo_id = :repo_id AND branches @> CAST(:branch_arr AS text[])", + {"repo_id": repo_id, "branch_arr": ["main"]}, + ) + a_nodes = list(_plan_nodes(plan_a)) + assert any(n.get("Index Name") == "ix_files_branches_gin" for n in a_nodes), a_nodes + assert not any(n.get("Node Type") == "Seq Scan" for n in a_nodes), a_nodes + assert not _projects_column(plan_a.get("Output"), "content"), plan_a + + plan_b = _explain( + conn, + "SELECT path, content_sha FROM files WHERE repo_id = :repo_id", + {"repo_id": repo_id}, + analyze=True, + ) + b_nodes = list(_plan_nodes(plan_b)) + assert not any(n.get("Node Type") == "Seq Scan" for n in b_nodes), b_nodes + io_scan = next((n for n in b_nodes if n.get("Index Name") == "uq_files_repo_path_sha"), None) + assert io_scan is not None, b_nodes + assert io_scan.get("Node Type") == "Index Only Scan", io_scan + assert io_scan.get("Heap Fetches") == 0, io_scan + # The two expensive mistakes read_repo_content_shas' docstring names: the + # Index Only Scan's own output list proves neither `content` nor `branches` + # is ever fetched -- content_sha/path/repo_id are the constraint's own + # columns, so an Index Only Scan is fundamentally incapable of returning + # anything else. + assert not _projects_column(io_scan.get("Output"), "content"), io_scan + assert not _projects_column(io_scan.get("Output"), "branches"), io_scan diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index d4ca8dd..de21d19 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -32,6 +32,7 @@ from app.config import Settings from app.db.models import INDEX_SEMANTICS_VERSION +from indexer.hashing import content_sha from indexer.job import ( BranchOutcome, RepoOutcome, @@ -325,6 +326,18 @@ def _noop_removed_fn(conn: Any, *, desired_repos: Any) -> list[str]: return [] +def _noop_shas_fn(conn: Any, *, name: str, branch: str) -> tuple[set[Any], set[Any]]: + """Default ``shas_fn`` for ``_run()`` -- ``_FakeConn`` cannot answer the real + ``read_indexed_shas`` query (it routes only on ``"repo_branches" in str(stmt)``, + per its own docstring), so any test that reaches the advisory read without an + explicit override would call the REAL primitive against the fake engine and + fail. Returning two empty sets degrades to "everything is changed/new", + matching ``read_indexed_shas``' own behaviour for a never-indexed repo -- safe + in the correct direction and inert for every test that never exercises delta + embedding at all.""" + return set(), set() + + class _RecordingReconcile: """Fake ``reconcile_retired_fn``/``reconcile_removed_fn`` pair, explicitly opted into. @@ -373,6 +386,7 @@ def _run( engine: _FakeEngine | None = None, reconcile_retired_fn: Any = _noop_retired_fn, reconcile_removed_fn: Any = _noop_removed_fn, + shas_fn: Any = _noop_shas_fn, ) -> int: """Drive run() with a faked config read but a REAL resolve_repos. @@ -382,6 +396,11 @@ def _run( against ``_FakeEngine``/``_FakeConn`` (neither implements ``conn.begin()``). Tests that assert on reconciliation itself pass an explicit :class:`_RecordingReconcile`'s bound methods. + + ``shas_fn`` defaults to :func:`_noop_shas_fn` for the same reason: a + version-matching, sha-mismatching stamp with semantic indexing on would + otherwise reach the REAL ``read_indexed_shas`` against ``_FakeConn``, which + cannot answer it (see that fake's docstring). """ wc = _FakeWorkspaceClient("tok") engine = engine if engine is not None else _FakeEngine() @@ -402,6 +421,7 @@ def _run( config_loader=lambda _client, _path: config, reconcile_retired_fn=reconcile_retired_fn, reconcile_removed_fn=reconcile_removed_fn, + shas_fn=shas_fn, ) @@ -756,6 +776,326 @@ def test_unbuildable_embedder_does_not_abort_the_whole_run() -> None: assert idx.chunk_writer is None +# --- file-level delta indexing: the chunk_writer covered-set guard (#104) --- + + +@pytest.mark.unit +def test_chunk_writer_covered_guard_skips_an_uncovered_path( + caplog: pytest.LogCaptureFixture, +) -> None: + """T14: a path outside _precompute_chunk_writer's own file list -> WARNING, and + chunk_writer issues NO statement at all (specifically no DELETE FROM chunks, + which write_chunks always opens with -- see indexer/chunk_store.py). This is + the defence-in-depth guard for a path the single-writer-per-repo invariant + says index_repo can never actually pass it; unreachable in production, but + the alternative (silently deleting an uncovered file's chunk rows) is worse + than a loud skip. + """ + from indexer.job import _precompute_chunk_writer + + pf = ParsedFile(path="ghost.py", lang="python", size=10, content="x = 1\n") + chunk_writer = _precompute_chunk_writer([], lambda texts: [[0.0] for _ in texts], 100) + conn = _FakeChunkConn() + with caplog.at_level(logging.WARNING, logger="indexer.job"): + chunk_writer(conn, 1, 99, pf) + assert conn.calls == [] + assert any( + "no precomputed chunks for ghost.py" in r.getMessage() + for r in caplog.records + if r.name == "indexer.job" + ) + + +@pytest.mark.unit +def test_chunk_writer_covers_every_embedded_path_including_zero_chunk_files() -> None: + """The covered set is `set(per_file)`, not `set(by_path)` -- a file that embeds + to zero chunks (e.g. an empty file) is still COVERED, so its chunk_writer call + reaches write_chunks([]) (the delete-only, zero-row-insert shape) rather than + the warn-and-skip guard above.""" + from indexer.job import _precompute_chunk_writer + + empty_pf = ParsedFile(path="empty.py", lang="python", size=0, content="") + chunk_writer = _precompute_chunk_writer([empty_pf], lambda texts: [[0.0] for _ in texts], 100) + conn = _FakeChunkConn() + chunk_writer(conn, 1, 99, empty_pf) + # write_chunks always issues its DELETE even for zero chunks (see + # indexer/chunk_store.py) -- so a real (non-warning) statement was issued. + assert len(conn.calls) >= 1 + + +# --- file-level delta indexing: the advisory shas_fn seam (#104) ------------ +# job.py's copy of indexer.store.read_repo_content_shas is ADVISORY -- it only +# narrows what _precompute_chunk_writer embeds. index_repo's own read (behind +# index_fn, unexercised by _RecordingIndex) is the authoritative classification; +# these tests pin job.py's half of the contract only: when shas_fn is called, +# with what, and what that narrows the embedder's input to. _DEFAULT_FILES is +# {"main.py": b"def f():\n return 1\n", "README.md": b"# hi\n"} (both text, +# both chunked -- indexer.parse.iter_chunks has no language gate). + +_MAIN_SHA = content_sha("def f():\n return 1\n") +_README_SHA = content_sha("# hi\n") + + +class _RecordingShas: + """Fake ``shas_fn``: records every ``(name, branch)`` call and returns a + scripted ``(carried, present)`` pair (``present`` is unused by job.py, which + only classifies "unchanged" from ``carried`` -- the provenance gate that + would also need ``present`` is index_repo's alone).""" + + def __init__( + self, + *, + carried: set[tuple[str, str]] | None = None, + present: set[tuple[str, str]] | None = None, + ) -> None: + self.calls: list[tuple[str, str]] = [] + self._carried = carried or set() + self._present = present or set() + + def __call__(self, conn: Any, *, name: str, branch: str) -> tuple[set[Any], set[Any]]: + self.calls.append((name, branch)) + return set(self._carried), set(self._present) + + +@pytest.mark.unit +def test_shas_fn_called_once_before_embedding_when_version_matches_and_sha_differs() -> None: + """T9: the gate is (stored version == current) AND (stored sha != HEAD) -- the + latter is implied by reaching this code at all (an exact stamp match skips the + branch entirely before any of this runs, see the Step 4 tests above).""" + order: list[str] = [] + shas = _RecordingShas() + + def _shas_fn(conn: Any, *, name: str, branch: str) -> tuple[set[Any], set[Any]]: + order.append("shas") + return shas(conn, name=name, branch=branch) + + def _embed(texts: list[str]) -> list[list[float]]: + order.append("embed") + return [[0.0] for _ in texts] + + engine = _FakeEngine(stamps={("acme/widgets", "main"): ("sha_old", INDEX_SEMANTICS_VERSION)}) + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + code = _run( + _config(repos=["acme/widgets"]), + _RecordingIndex(), + cfg=cfg, + embed_fn=_embed, + engine=engine, + shas_fn=_shas_fn, + ) + assert code == 0 + assert shas.calls == [("acme/widgets", "main")] + assert order == ["shas", "embed"] + + +@pytest.mark.unit +def test_shas_fn_not_called_when_the_semantics_version_is_stale() -> None: + """T10: version mismatch -> the delta gate is closed, so job.py never issues the + advisory read either, and the embedder receives every file's chunk text.""" + shas = _RecordingShas() + embed_calls: list[list[str]] = [] + + def _embed(texts: list[str]) -> list[list[float]]: + embed_calls.append(list(texts)) + return [[0.0] for _ in texts] + + engine = _FakeEngine( + stamps={("acme/widgets", "main"): ("sha_widgets", INDEX_SEMANTICS_VERSION - 1)} + ) + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + code = _run( + _config(repos=["acme/widgets"]), + _RecordingIndex(), + cfg=cfg, + embed_fn=_embed, + engine=engine, + shas_fn=shas, + ) + assert code == 0 + assert shas.calls == [] + assert len(embed_calls) == 1 + # Both main.py's and README.md's chunk text are present -- nothing narrowed. + joined = "\n".join(embed_calls[0]) + assert "def f" in joined + assert "# hi" in joined + + +@pytest.mark.unit +def test_shas_fn_not_called_for_a_never_indexed_repo() -> None: + """T10b: a missing stamp degrades to (None, None) -- version is None, never equal + to INDEX_SEMANTICS_VERSION, so this is the same gate-closed path as a stale + version, exercised separately because it is the far more common real case (a + repo's first index) than an explicit version regression.""" + shas = _RecordingShas() + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + code = _run( + _config(repos=["acme/widgets"]), + _RecordingIndex(), + cfg=cfg, + embed_fn=lambda texts: [[0.0] for _ in texts], + shas_fn=shas, + ) + assert code == 0 + assert shas.calls == [] + + +@pytest.mark.unit +def test_shas_fn_narrows_the_embed_set_to_not_unchanged_files() -> None: + """T11: README.md is reported unchanged (carried); main.py is not -- the embedder + must receive only main.py's chunk text.""" + embed_calls: list[list[str]] = [] + + def _embed(texts: list[str]) -> list[list[float]]: + embed_calls.append(list(texts)) + return [[0.0] for _ in texts] + + shas_fn = _RecordingShas(carried={("README.md", _README_SHA)}) + engine = _FakeEngine(stamps={("acme/widgets", "main"): ("sha_old", INDEX_SEMANTICS_VERSION)}) + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + code = _run( + _config(repos=["acme/widgets"]), + _RecordingIndex(), + cfg=cfg, + embed_fn=_embed, + engine=engine, + shas_fn=shas_fn, + ) + assert code == 0 + assert len(embed_calls) == 1 + joined = "\n".join(embed_calls[0]) + assert "def f" in joined + assert "# hi" not in joined + + +@pytest.mark.unit +def test_shas_fn_all_unchanged_calls_the_embedder_with_zero_texts() -> None: + """T12 (issue AC1): every file carried -> the embed list is empty, and + _precompute_chunk_writer's own `all_texts` guard means the embedder is not + even called (matching read_indexed_shas' safe-degrade convention: an unused + injection point is never exercised, not called with an empty list).""" + embed_calls: list[list[str]] = [] + + def _embed(texts: list[str]) -> list[list[float]]: + embed_calls.append(list(texts)) + return [[0.0] for _ in texts] + + shas_fn = _RecordingShas( + carried={("README.md", _README_SHA), ("main.py", _MAIN_SHA)}, + ) + engine = _FakeEngine(stamps={("acme/widgets", "main"): ("sha_old", INDEX_SEMANTICS_VERSION)}) + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + idx = _RecordingIndex() + code = _run( + _config(repos=["acme/widgets"]), + idx, + cfg=cfg, + embed_fn=_embed, + engine=engine, + shas_fn=shas_fn, + ) + assert code == 0 + assert embed_calls == [] + # The core index still ran over BOTH files -- narrowing the embed set never + # narrows what index_fn (the store) sees; that is index_repo's classification + # to make, from its own authoritative read. + assert idx.counts == [IndexCounts(files=2, symbols=1, swept=0, edges=0)] + assert idx.chunk_writer is not None + + +@pytest.mark.unit +def test_indexed_summary_line_is_byte_identical_regardless_of_delta_narrowing( + caplog: pytest.LogCaptureFixture, +) -> None: + """T13: the drain loop's `indexed name@branch: files=.. symbols=.. edges=.. swept=..` + line is IndexCounts' own format -- narrowing the embed set must not touch it.""" + shas_fn = _RecordingShas(carried={("README.md", _README_SHA)}) + engine = _FakeEngine(stamps={("acme/widgets", "main"): ("sha_old", INDEX_SEMANTICS_VERSION)}) + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + _config(repos=["acme/widgets"]), + _RecordingIndex(), + cfg=cfg, + embed_fn=lambda texts: [[0.0] for _ in texts], + engine=engine, + shas_fn=shas_fn, + ) + assert code == 0 + indexed_lines = [ + r.getMessage() + for r in caplog.records + if r.name == "indexer.job" and r.getMessage().startswith("indexed ") + ] + assert indexed_lines == ["indexed acme/widgets@main: files=2 symbols=1 edges=0 swept=0"] + + +# --- file-level delta indexing: degraded-semantics run summary (#104) ------- + + +@pytest.mark.unit +def test_precompute_failure_marks_the_branch_outcome_degraded() -> None: + """A chunk-precompute failure marks the resulting BranchOutcome, not just the + per-branch WARNING already covered by + test_embedder_failure_degrades_but_still_indexes_the_core.""" + from indexer.job import _index_one_branch + + def _down(_texts: list[str]) -> list[list[float]]: + raise RuntimeError("serving endpoint unavailable") + + with httpx.Client(transport=httpx.MockTransport(_GitHub())) as client: + outcome = _index_one_branch( + "acme/widgets", + org="acme", + repo="widgets", + branch="main", + is_default=True, + default_head_sha="sha_widgets", + http_client=client, + engine=_FakeEngine(), + index_fn=_RecordingIndex(), + cfg=Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100), + embed_fn=_down, + stamps={}, + started=time.monotonic(), + max_chunks_per_repo=100, + shas_fn=_noop_shas_fn, + ) + assert outcome.status == "indexed" + assert outcome.semantic_degraded is True + + +@pytest.mark.unit +def test_run_emits_one_aggregate_warning_naming_every_degraded_branch( + caplog: pytest.LogCaptureFixture, +) -> None: + """The run-completion WARNING is the greppable record the runbook remedy + (clear the branch's semantics stamp) depends on -- the per-branch warning at + the precompute site is easy to miss in a large run's log.""" + + def _down(_texts: list[str]) -> list[list[float]]: + raise RuntimeError("serving endpoint unavailable") + + cfg = Settings(semantic_enabled=True) + with caplog.at_level(logging.WARNING, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), _RecordingIndex(), cfg=cfg, embed_fn=_down) + assert code == 0 + warnings = [r.getMessage() for r in caplog.records if r.name == "indexer.job"] + aggregate = [m for m in warnings if m.startswith("1 branch(es) finished with degraded")] + assert len(aggregate) == 1 + assert "acme/widgets@main" in aggregate[0] + + +@pytest.mark.unit +def test_run_emits_no_aggregate_warning_when_nothing_degraded( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger="indexer.job"): + code = _run(_config(repos=["acme/widgets"]), _RecordingIndex()) + assert code == 0 + warnings = [r.getMessage() for r in caplog.records if r.name == "indexer.job"] + assert not any("degraded semantic coverage" in m for m in warnings) + + # --- config.yaml `semantic:` overlay onto cfg (config.yaml > env > default) -- @@ -1224,6 +1564,7 @@ def test_index_one_inner_reports_discovery_complete_for_a_normal_run() -> None: cfg=Settings(semantic_enabled=False), embed_fn=None, stamps={}, + shas_fn=_noop_shas_fn, ) assert outcome.name == "acme/widgets" assert outcome.discovery_complete is True @@ -1250,6 +1591,7 @@ def test_index_one_inner_reports_discovery_incomplete_when_capped() -> None: cfg=Settings(semantic_enabled=False), embed_fn=None, stamps={}, + shas_fn=_noop_shas_fn, ) assert outcome.discovery_complete is False assert len(outcome.outcomes) == SOFT_BRANCH_CAP @@ -1277,6 +1619,7 @@ def test_index_one_inner_default_flip_mirror_is_complete() -> None: cfg=Settings(semantic_enabled=False), embed_fn=None, stamps={}, + shas_fn=_noop_shas_fn, ) assert [o.branch for o in outcome.outcomes] == ["main"] assert outcome.discovery_complete is True @@ -1584,6 +1927,7 @@ def test_repo_context_is_reset_even_when_the_repo_fails() -> None: cfg=Settings(semantic_enabled=False), embed_fn=None, stamps={}, + shas_fn=_noop_shas_fn, ) assert _repo_ctx.get() == "-" diff --git a/tests/unit/test_store_delta.py b/tests/unit/test_store_delta.py new file mode 100644 index 0000000..4f90ff0 --- /dev/null +++ b/tests/unit/test_store_delta.py @@ -0,0 +1,523 @@ +"""Unit tests for ``index_repo``'s file-level delta path, at the STATEMENT level. + +A hand-rolled fake ``Connection`` (the ``_FakeConn`` idiom from +``tests/unit/test_store_chunk_writer.py``, extended to answer the two projection +reads and the provenance gate) stands in for Postgres, so these are true unit +tests: what they pin is the exact *statement inventory* each classification +produces, and the order the transaction issues them in. Row-identity proof -- +that skipping really does leave the stored serials untouched -- is +``tests/integration/test_store_delta.py``'s job, against real SQL. + +The fake records one short label per executed statement (``repos-insert``, +``read-carried``, ``file-upsert``, ...) rather than the statement objects, so an +assertion reads as the inventory it is checking. +""" + +from __future__ import annotations + +import contextlib +import logging +from typing import Any, NamedTuple + +import pytest +from sqlalchemy import Delete, Insert, Update + +from app.db.models import INDEX_SEMANTICS_VERSION +from indexer.hashing import content_sha +from indexer.languages import ( + ExtractedEdge, + ExtractedSymbol, + FileExtraction, + IndexCounts, + ParsedFile, +) +from indexer.store import index_repo + + +class _ShaRow(NamedTuple): + """The exact projection statements 3a/3b return -- path and content_sha, nothing else.""" + + path: str + content_sha: str + + +class _IdRow(NamedTuple): + """The batched membership ``UPDATE ... RETURNING id, path, content_sha`` shape.""" + + id: int + path: str + content_sha: str + + +class _FakeResult: + def __init__( + self, + *, + scalar: Any = None, + rowcount: int = 0, + row: Any = None, + rows: list[Any] | None = None, + ) -> None: + self._scalar = scalar + self._row = row + self._rows = rows or [] + self.rowcount = rowcount + + def scalar_one(self) -> Any: + return self._scalar + + def scalar_one_or_none(self) -> Any: + return self._scalar + + def one(self) -> Any: + return self._row + + def all(self) -> list[Any]: + return self._rows + + def __iter__(self) -> Any: + return iter(self._rows) + + +class _FakeConn: + """Answers every statement ``index_repo`` can issue, and labels it. + + ``baseline`` is what statement 2's ``RETURNING`` yields -- the pair that + opens or closes the delta gate. ``carried``/``present`` script statements 3a + and 3b (see ``indexer.store.read_repo_content_shas``); ``provenance`` + scripts statement 4's ``NOT EXISTS`` gate. + """ + + def __init__( + self, + *, + baseline: tuple[str | None, int | None] = (None, None), + carried: set[tuple[str, str]] | None = None, + present: set[tuple[str, str]] | None = None, + provenance: bool = True, + stamp_rowcount: int = 1, + ) -> None: + self._next_file_id = 1 + self._baseline = baseline + self._carried = sorted(carried or set()) + self._present = sorted(present or set()) + self._provenance = provenance + self._stamp_rowcount = stamp_rowcount + self.kinds: list[str] = [] + self.stamp_values: dict[str, Any] = {} + self.membership_params: dict[str, Any] = {} + + def begin(self) -> Any: + return contextlib.nullcontext() + + def _text_execute(self, sql: str, params: Any) -> _FakeResult: + # Ordered most-specific-first: statement 3b's text is a PREFIX of 3a's. + if "SELECT path, content_sha FROM files" in sql and "branches @>" in sql: + self.kinds.append("read-carried") + return _FakeResult(rows=[_ShaRow(p, s) for p, s in self._carried]) + if "SELECT path, content_sha FROM files" in sql: + self.kinds.append("read-present") + return _FakeResult(rows=[_ShaRow(p, s) for p, s in self._present]) + if "SELECT NOT EXISTS" in sql: + self.kinds.append("provenance-gate") + return _FakeResult(scalar=self._provenance) + if "UPDATE files" in sql and "array_agg(DISTINCT e)" in sql: + self.kinds.append("membership-union") + self.membership_params = dict(params or {}) + rows = [] + for path, sha in zip( + (params or {}).get("paths", []), (params or {}).get("shas", []), strict=True + ): + rows.append(_IdRow(self._next_file_id, path, sha)) + self._next_file_id += 1 + return _FakeResult(rows=rows, rowcount=len(rows)) + if "UPDATE files SET branches = array_remove" in sql: + self.kinds.append("sweep-update") + return _FakeResult(rowcount=0) + if "DELETE FROM files" in sql: + self.kinds.append("sweep-delete") + return _FakeResult(rowcount=0) + raise AssertionError(f"unexpected text() statement: {sql!r}") + + def execute(self, stmt: Any, params: Any = None) -> _FakeResult: + sql = getattr(stmt, "text", None) + if sql is not None: + return self._text_execute(sql, params) + + table = stmt.table.name + if isinstance(stmt, Insert) and table == "repos": + self.kinds.append("repos-insert") + return _FakeResult(scalar=1) + if isinstance(stmt, Insert) and table == "repo_branches": + self.kinds.append("repo-branches-insert") + return _FakeResult(row=self._baseline) + if isinstance(stmt, Insert) and table == "files": + self.kinds.append("file-upsert") + file_id = self._next_file_id + self._next_file_id += 1 + return _FakeResult(scalar=file_id) + if isinstance(stmt, Delete) and table == "symbols": + self.kinds.append("symbols-delete") + return _FakeResult() + if isinstance(stmt, Insert) and table == "symbols": + self.kinds.append("symbols-insert") + return _FakeResult() + if isinstance(stmt, Delete) and table == "reference_edges": + self.kinds.append("edges-delete") + return _FakeResult() + if isinstance(stmt, Insert) and table == "reference_edges": + self.kinds.append("edges-insert") + return _FakeResult() + if isinstance(stmt, Update) and table == "repo_branches": + self.kinds.append("stamp") + self.stamp_values = dict(stmt.compile().params) + return _FakeResult(rowcount=self._stamp_rowcount) + raise AssertionError(f"unexpected statement against {table!r}: {stmt}") + + +def _pf(path: str, content: str) -> ParsedFile: + return ParsedFile(path=path, lang="python", size=len(content.encode()), content=content) + + +def _item( + path: str, content: str, *, symbols: bool = True, edges: bool = False +) -> tuple[ParsedFile, FileExtraction]: + symbol = ExtractedSymbol("f", "function", 1, 2) + return ( + _pf(path, content), + FileExtraction( + symbols=[symbol] if symbols else [], + edges=[ExtractedEdge(kind="call", target="t", line=2, enclosing=symbol)] + if edges + else [], + ), + ) + + +def _key(path: str, content: str) -> tuple[str, str]: + return (path, content_sha(content)) + + +def _index(conn: _FakeConn, items: Any, **kwargs: Any) -> IndexCounts: + return index_repo( + conn, + name="acme/widgets", + branch="main", + is_default=True, + head_sha="sha_new", + items=items, + **kwargs, + ) + + +# --- The gate: closed unless the stored version equals the current one -------- + + +@pytest.mark.unit +@pytest.mark.parametrize( + "baseline", + [(None, None), ("sha_old", None), ("sha_old", INDEX_SEMANTICS_VERSION - 1)], + ids=["never-indexed", "null-version", "older-version"], +) +def test_version_mismatch_never_issues_the_projection_reads( + baseline: tuple[str | None, int | None], +) -> None: + """T1: a NULL or stale stored version means full path for every file, and the + two projection reads (and the provenance gate) are never issued at all.""" + conn = _FakeConn(baseline=baseline, carried={_key("a.py", "x = 1\n")}) + counts = _index(conn, [_item("a.py", "x = 1\n")]) + + assert "read-carried" not in conn.kinds + assert "read-present" not in conn.kinds + assert "provenance-gate" not in conn.kinds + assert "file-upsert" in conn.kinds + assert counts == IndexCounts(files=1, symbols=1, swept=0, edges=0) + + +@pytest.mark.unit +def test_delta_on_all_unchanged_writes_nothing() -> None: + """T2 (AC1): every parsed file already carried by this branch at this content + means zero files/symbols/reference_edges statements -- but the sweep still + runs against the FULL seen-set, and the stamp is still last.""" + items = [_item("a.py", "x = 1\n"), _item("b.py", "y = 2\n")] + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("a.py", "x = 1\n"), _key("b.py", "y = 2\n")}, + present={_key("a.py", "x = 1\n"), _key("b.py", "y = 2\n")}, + ) + counts = _index(conn, items) + + assert conn.kinds == [ + "repos-insert", + "repo-branches-insert", + "read-carried", + "read-present", + "provenance-gate", + "sweep-update", + "sweep-delete", + "stamp", + ] + # files still counts files SEEN this run (the seen-set size the sweep uses), + # so the sweep and its empty-seen-set guard need no delta awareness. + assert counts == IndexCounts(files=2, symbols=0, swept=0, edges=0) + + +@pytest.mark.unit +def test_all_unchanged_run_calls_no_chunk_writer() -> None: + """T2 (AC1), chunk half: an unchanged file's chunk rows are never rewritten.""" + calls: list[str] = [] + items = [_item("a.py", "x = 1\n")] + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("a.py", "x = 1\n")}, + present={_key("a.py", "x = 1\n")}, + ) + _index( + conn, + items, + chunk_writer=lambda _c, _r, _f, pf: calls.append(pf.path), + ) + assert calls == [] + + +@pytest.mark.unit +def test_changed_content_at_a_known_path_takes_the_full_path() -> None: + """A path whose content moved is NOT in ``carried`` under its new sha, so it + takes the full write path -- the (path, content_sha) keying, not path alone.""" + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("a.py", "x = 1\n")}, + present={_key("a.py", "x = 1\n")}, + ) + counts = _index(conn, [_item("a.py", "x = 999\n", edges=True)]) + + assert conn.kinds.count("file-upsert") == 1 + assert conn.kinds.count("symbols-delete") == 1 + assert conn.kinds.count("edges-delete") == 1 + assert counts == IndexCounts(files=1, symbols=1, swept=0, edges=1) + + +# --- Membership-only: one batched union, a chunk write, no symbol/edge work --- + + +@pytest.mark.unit +def test_membership_only_issues_one_batched_union_and_no_symbol_work() -> None: + """T3 (AC3): a file stored for this repo but not carried by this branch takes + the membership path -- ONE batched UPDATE for the whole class, one + chunk_writer call per file, and no symbols/reference_edges statements.""" + calls: list[tuple[int, str]] = [] + items = [_item("a.py", "x = 1\n"), _item("b.py", "y = 2\n")] + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried=set(), + present={_key("a.py", "x = 1\n"), _key("b.py", "y = 2\n")}, + ) + counts = _index(conn, items, chunk_writer=lambda _c, _r, fid, pf: calls.append((fid, pf.path))) + + assert conn.kinds.count("membership-union") == 1 + assert "file-upsert" not in conn.kinds + assert "symbols-delete" not in conn.kinds + assert "edges-delete" not in conn.kinds + # The file_id each chunk write used came from the UPDATE's RETURNING, not a + # second lookup. + assert calls == [(1, "a.py"), (2, "b.py")] + assert conn.membership_params["paths"] == ["a.py", "b.py"] + assert conn.membership_params["branch_arr"] == ["main"] + # symbols/edges legitimately fall to zero: no rows were inserted. + assert counts == IndexCounts(files=2, symbols=0, swept=0, edges=0) + + +@pytest.mark.unit +def test_membership_only_is_refused_when_a_sibling_branch_is_stale() -> None: + """T4 (AC6): statement 4 false -- some branch of this repo sits at another + semantics version -- forces every would-be membership file down the full + path, so it can never inherit stale-version symbols/edges.""" + items = [_item("a.py", "x = 1\n")] + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried=set(), + present={_key("a.py", "x = 1\n")}, + provenance=False, + ) + counts = _index(conn, items) + + assert "provenance-gate" in conn.kinds + assert "membership-union" not in conn.kinds + assert conn.kinds.count("file-upsert") == 1 + assert conn.kinds.count("symbols-delete") == 1 + assert counts == IndexCounts(files=1, symbols=1, swept=0, edges=0) + + +@pytest.mark.unit +def test_mixed_classification_statement_inventory() -> None: + """T5 (AC2): 1 unchanged, 1 membership-only, 1 changed, 1 new -> the exact + inventory, with the batched union issued once, AFTER the per-file loop.""" + unchanged = _item("keep.py", "k = 1\n") + member = _item("shared.py", "s = 1\n") + changed = _item("moved.py", "m = 2\n") + added = _item("new.py", "n = 1\n") + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("keep.py", "k = 1\n"), _key("moved.py", "m = 1\n")}, + present={ + _key("keep.py", "k = 1\n"), + _key("moved.py", "m = 1\n"), + _key("shared.py", "s = 1\n"), + }, + ) + counts = _index(conn, [unchanged, member, changed, added]) + + assert conn.kinds == [ + "repos-insert", + "repo-branches-insert", + "read-carried", + "read-present", + "provenance-gate", + # moved.py -- changed content at a known path + "file-upsert", + "symbols-delete", + "symbols-insert", + "edges-delete", + # new.py -- never seen + "file-upsert", + "symbols-delete", + "symbols-insert", + "edges-delete", + # shared.py -- the whole membership class, batched, after the loop + "membership-union", + "sweep-update", + "sweep-delete", + "stamp", + ] + assert conn.membership_params["paths"] == ["shared.py"] + assert counts == IndexCounts(files=4, symbols=2, swept=0, edges=0) + + +@pytest.mark.unit +def test_empty_membership_class_issues_no_union_statement() -> None: + """T5, the stability half: an empty membership set is SKIPPED, never issued + as a no-op UPDATE -- which is what keeps the inventories above stable.""" + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("a.py", "x = 1\n")}, + present={_key("a.py", "x = 1\n")}, + ) + _index(conn, [_item("a.py", "x = 1\n")]) + assert "membership-union" not in conn.kinds + + +@pytest.mark.unit +def test_membership_without_a_chunk_writer_issues_only_the_union() -> None: + """The semantic-off path: the union still runs, nothing else does.""" + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried=set(), + present={_key("a.py", "x = 1\n")}, + ) + _index(conn, [_item("a.py", "x = 1\n")], chunk_writer=None) + assert conn.kinds.count("membership-union") == 1 + + +# --- The `delta write set` line: always emitted, one shape ------------------- + + +@pytest.mark.unit +def test_delta_write_set_line_reports_the_breakdown_with_the_gate_open( + caplog: pytest.LogCaptureFixture, +) -> None: + """T7: the breakdown rides its OWN line from indexer.store -- IndexCounts is + unchanged, so this is where unchanged/membership become visible.""" + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("keep.py", "k = 1\n")}, + present={_key("keep.py", "k = 1\n"), _key("shared.py", "s = 1\n")}, + ) + with caplog.at_level(logging.INFO, logger="indexer.store"): + _index( + conn, + [ + _item("keep.py", "k = 1\n"), + _item("shared.py", "s = 1\n"), + _item("new.py", "n = 1\n"), + ], + ) + + assert ( + "acme/widgets@main: delta write set 1/3 files " + "(unchanged=1 membership=1, semantics gate open)" + ) in caplog.text + + +@pytest.mark.unit +def test_delta_write_set_line_is_present_with_the_gate_closed( + caplog: pytest.LogCaptureFixture, +) -> None: + """T7, the other half: the line never disappears -- a closed gate reports the + full-path reason instead, so no grep an operator writes breaks.""" + conn = _FakeConn(baseline=("sha_old", INDEX_SEMANTICS_VERSION - 1)) + with caplog.at_level(logging.INFO, logger="indexer.store"): + _index(conn, [_item("a.py", "x = 1\n"), _item("b.py", "y = 2\n")]) + + assert ( + f"acme/widgets@main: delta write set 2/2 files (unchanged=0 membership=0, " + f"semantics gate closed: stored v{INDEX_SEMANTICS_VERSION - 1} " + f"!= v{INDEX_SEMANTICS_VERSION})" + ) in caplog.text + + +# --- Transaction shape and the untouched guards ------------------------------ + + +@pytest.mark.unit +def test_transaction_shape_is_pinned() -> None: + """T6: the epic's non-negotiable rule. repos first, repo_branches second, the + projection reads third, the CAS stamp LAST -- whatever the classification.""" + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("a.py", "x = 1\n")}, + present={_key("a.py", "x = 1\n"), _key("b.py", "y = 2\n")}, + ) + _index(conn, [_item("a.py", "x = 1\n"), _item("c.py", "z = 3\n")]) + + assert conn.kinds[0] == "repos-insert" + assert conn.kinds[1] == "repo-branches-insert" + assert conn.kinds[2] == "read-carried" + assert conn.kinds[3] == "read-present" + assert conn.kinds[4] == "provenance-gate" + assert conn.kinds[-1] == "stamp" + + +@pytest.mark.unit +def test_empty_items_still_skips_the_sweep_with_delta_on() -> None: + """T8: the empty-seen-set guard is untouched by the delta path.""" + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("a.py", "x = 1\n")}, + present={_key("a.py", "x = 1\n")}, + ) + counts = _index(conn, []) + + assert "sweep-update" not in conn.kinds + assert "sweep-delete" not in conn.kinds + assert counts == IndexCounts(files=0, symbols=0, swept=0, edges=0) + + +@pytest.mark.unit +def test_zero_parse_run_does_not_advance_the_semantics_version() -> None: + """The §2.3 base case, at the statement level: an empty seen-set stamps the + commit but leaves ``index_semantics_version`` at the statement-2 baseline.""" + conn = _FakeConn(baseline=("sha_old", INDEX_SEMANTICS_VERSION - 1)) + _index(conn, []) + + assert conn.stamp_values["last_indexed_commit"] == "sha_new" + assert conn.stamp_values["index_semantics_version"] == INDEX_SEMANTICS_VERSION - 1 + + +@pytest.mark.unit +def test_non_empty_run_advances_the_semantics_version() -> None: + """The complement: a run that wrote something DOES claim the current version.""" + conn = _FakeConn(baseline=("sha_old", INDEX_SEMANTICS_VERSION - 1)) + _index(conn, [_item("a.py", "x = 1\n")]) + + assert conn.stamp_values["last_indexed_commit"] == "sha_new" + assert conn.stamp_values["index_semantics_version"] == INDEX_SEMANTICS_VERSION From ad04748651d3de530dcc79a4c96d4da8231c5dc9 Mon Sep 17 00:00:00 2001 From: Tanner <84605639+IceRhymers@users.noreply.github.com> Date: Sat, 25 Jul 2026 09:32:12 -0700 Subject: [PATCH 4/9] indexer: single-pass in-memory tarball ingestion (#106) Refs #106 --- config.yaml | 9 +- docs/runbooks/indexing-parallelism.md | 83 ++- indexer/AGENTS.md | 17 +- indexer/fetch.py | 66 +- indexer/ingest.py | 365 +++++++++++ indexer/job.py | 42 +- indexer/repo_config.py | 13 +- indexer/timing.py | 4 +- tests/integration/test_job_ingest_delta.py | 147 +++++ tests/unit/test_fetch.py | 68 +-- tests/unit/test_ingest.py | 678 +++++++++++++++++++++ tests/unit/test_ingest_parity.py | 145 +++++ tests/unit/test_job.py | 193 +++++- 13 files changed, 1640 insertions(+), 190 deletions(-) create mode 100644 indexer/ingest.py create mode 100644 tests/integration/test_job_ingest_delta.py create mode 100644 tests/unit/test_ingest.py create mode 100644 tests/unit/test_ingest_parity.py diff --git a/config.yaml b/config.yaml index 2ad8ec6..6094bfa 100644 --- a/config.yaml +++ b/config.yaml @@ -4,12 +4,13 @@ version: 1 # How many repos the job indexes at once. Default 4, max 8, min 1. -# This is a DISK bound, not a CPU one: each worker holds its 500 MB tarball and -# its 2 GB extraction alive at the same time, so budget 2.5 GB per worker — -# 10 GB at the default 4, 20 GB at the ceiling of 8. +# This is a DISK bound, not a CPU one: each worker holds its 500 MB tarball on +# disk. That is the only on-disk artifact — the archive is streamed in memory and +# never extracted — so budget 0.5 GB per worker: 2 GB at the default 4, 4 GB at +# the ceiling of 8. # Returns at the ceiling are sublinear: symbol extraction does not parallelise # (measured 0.95x on 4 threads), so you buy far less than 8x for a hard linear -# 20 GB of disk. Raise it knowing that. +# 4 GB of disk. Raise it knowing that. # With semantic indexing enabled this is clamped to 2 — a MEMORY bound, since # embedding materialises a whole repo's chunks (~0.5-0.8 GB per worker). # index_concurrency: 4 diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index 9de234b..892e0c3 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/docs/runbooks/indexing-parallelism.md @@ -114,10 +114,10 @@ of their own. Records emitted outside a worker (config resolution, the drain loop, third-party libraries) carry `-`. ``` -INFO indexer.job [-]: local disk at /tmp: 41.2 GB free of 64.0 GB total; 4 worker(s) x 2.5 GB peak +INFO indexer.job [-]: local disk at /tmp: 41.2 GB free of 64.0 GB total; 4 worker(s) x 0.5 GB peak INFO indexer.fetch [acme/widgets]: ... -INFO indexer.job [acme/widgets]: phase timing acme/widgets@main: total=213.32s resolve=0.00s download=12.10s extract=8.40s parse=31.00s embed=88.20s db=64.50s sweep=0.30s other=8.82s -INFO indexer.job [acme/widgets]: finished acme/widgets in 213.74s (resolve=0.42s list=0.00s) +INFO indexer.job [acme/widgets]: phase timing acme/widgets@main: total=204.92s resolve=0.00s download=12.10s parse=31.00s embed=88.20s db=64.50s sweep=0.30s other=8.82s +INFO indexer.job [acme/widgets]: finished acme/widgets in 205.34s (resolve=0.42s list=0.00s) INFO indexer.job [acme/gadgets]: skipped acme/gadgets@main: already indexed at abc123 (semantics v1) in 0.41s ``` @@ -141,15 +141,21 @@ to attribute. grep 'phase timing' run.log # one line per indexed branch ``` -The nine fields are fixed, always present, always in this order, always `%.2fs`. +The eight fields are fixed, always present, always in this order, always `%.2fs`. A phase that did not run prints `0.00s` rather than disappearing, so the line never changes shape between a semantic-on and a semantic-off run and every grep you write keeps working. Read the largest field; that is the branch's bottleneck. +(There used to be a ninth, `extract=`. #106 removed the phase itself — the +tarball is streamed once, in memory, and never extracted — so the field was +deleted rather than pinned at `0.00s`. The "prints `0.00s` rather than +disappearing" rule is about one build's semantic-on vs semantic-off runs, not a +promise that the field set never changes across releases.) + | Dominant phase | What it means | Which issue addresses it | |---|---|---| -| `download` / `extract` | archive I/O bound | #106 (single-pass in-memory ingestion) | -| `parse` | GIL-bound tree-sitter extraction | #108 (process-pool extraction) | +| `download` | archive I/O bound | — (the decompression it used to be paired with is now fused into `parse`, #106) | +| `parse` | GIL-bound tree-sitter extraction — **plus, since #106, the archive's gzip decompression, tar-stream read and UTF-8 decode**, which used to be the separate `extract=` field | #108 (process-pool extraction) addresses the tree-sitter half only | | `embed` | serial AI Gateway round trips | #107 (concurrent embedding) | | `db` | per-file round trips | #105 (batched writes) | | any of the above, on **unchanged** content | redundant work | #104 (file-level delta indexing) | @@ -158,7 +164,7 @@ you write keeps working. Read the largest field; that is the branch's bottleneck its size — but it does NOT touch `parse`: extraction still runs on every file every run (tree-sitter must produce a `FileExtraction` before `index_repo` can classify it), so an all-unchanged branch on a large repo -still pays its full `download`+`extract`+`parse` cost. See `indexer.store`'s +still pays its full `download`+`parse` cost. See `indexer.store`'s `delta write set …` line (below) to tell "this branch is genuinely mostly-new" from "this branch is mostly-unchanged but still parsing everything" — the latter is exactly the case `#108` (process-pool extraction) or a future @@ -176,10 +182,13 @@ Four fields need interpretation before you act on them: `finished … in Xs` is measured on its own clock and is deliberately not reconciled against the parenthesised numbers. - **`other=` is the unattributed residual**, `total` minus every measured phase, - clamped at zero. It is dominated by the temp-dir teardown — an `rm -rf` of a - freshly extracted multi-GB tree — plus the pre-flight disk check. It exists so - the line has no silently missing time; a large `other` means something real is - happening outside every instrumented phase and is worth chasing. + clamped at zero. It covers the temp-dir teardown plus the pre-flight disk + check. Since #106 that teardown is an `rm -rf` of one compressed tarball, not + of a freshly extracted multi-GB tree, so `other` shrinks materially on large + repos — if you are reading an old run's numbers, do not go hunting for a + teardown cost that no longer exists. It exists so the line has no silently + missing time; a large `other` means something real is happening outside every + instrumented phase and is worth chasing. - **`embed=` covers chunking as well as the network.** It spans `iter_chunks` (CPU/GIL-bound) *and* the serial AI Gateway round trips. #107 addresses only the round trips, so before routing work there, confirm the phase is @@ -189,10 +198,19 @@ Four fields need interpretation before you act on them: transaction.** Files stream lazily through `index_repo`'s open transaction for bounded memory, so file production is timed separately and subtracted from `db`; the sweep is subtracted too. On the **non-semantic** path, though, the - directory walk itself materializes inside that open transaction (`parse.py`'s - `rglob`, on the first item). That is long-standing behavior which this - instrumentation merely makes visible for the first time — it is not a new - regression. + file walk itself materializes inside that open transaction — since #106 that + is the tar stream (`ingest.py`), not `parse.py`'s `rglob`, on the first item. + That is long-standing behavior which this instrumentation merely makes visible + for the first time — it is not a new regression. What #106 *did* move into that + window is **archive validation**: the decompression-bomb cap, the + exactly-one-top-level-dir check, member-name and link-target safety, and any + `tarfile` corruption error are now raised as the stream is consumed rather than + before the connection is taken. A malformed archive therefore surfaces as a + rolled-back transaction and one briefly-held pooled connection instead of a + pre-connection failure. **The branch-level outcome is unchanged** — + `failed`, exit code non-zero, nothing written. On the semantic (production) + path nothing moved at all: the file list is materialized up front, so the + archive is fully validated before any connection is opened. ### 2.2 The delta write set line (#104) @@ -232,25 +250,30 @@ provenance gate). |---|---|---| | `index_concurrency` | 1..8, default **4** | Repos in flight | | `MAX_TARBALL_BYTES` | 500 MB | The compressed download, per worker | -| `MAX_EXTRACTED_BYTES` | 2 GB | The uncompressed tree, per worker | +| `MAX_EXTRACTED_BYTES` | 2 GB | The streamed uncompressed content, per branch | -The two byte caps **sum**, they do not `max()`: the tarball stays on disk inside -the worker's temp directory while the extraction grows beside it. Peak local -disk is therefore: +Only the first of the two byte caps is a **disk** cap. Since #106 the tarball is +streamed once, in memory, and is never extracted, so `MAX_EXTRACTED_BYTES` is a +**work** cap — a decompression-bomb guard on how much content one branch may pull +out of its archive — and it lives in `indexer/ingest.py`, beside its only +consumer, rather than in `indexer/fetch.py`. The two therefore no longer sum: +the compressed tarball is the only artifact on disk, so peak local disk is +`index_concurrency` × 500 MB: | `index_concurrency` | Peak local disk | |---|---| -| 1 | 2.5 GB | -| 2 | 5 GB | -| **4 (default)** | **10 GB** | -| 8 (ceiling) | 20 GB | +| 1 | 0.5 GB | +| 2 | 1 GB | +| **4 (default)** | **2 GB** | +| 8 (ceiling) | 4 GB | **Returns at the ceiling are sublinear; the disk cost is not.** Symbol extraction was measured at **0.95x on 4 threads** — the tree walk is GIL-serialized and is ~56% of extraction time, so Amdahl's law caps the speedup -well below 8x. Meanwhile the 20 GB is a hard, linear, unavoidable cost. Raise +well below 8x. Meanwhile the 4 GB is a hard, linear, unavoidable cost. Raise `index_concurrency` to 8 only knowing you are buying a fraction of a speedup -with a doubling of disk. +with a doubling of disk. (#106 lowered these numbers by 5x but deliberately did +**not** move the default of 4; re-deriving it is #109's job.) **Semantic indexing clamps the pool to 2**, regardless of `index_concurrency`. That clamp is a *memory* bound, not a CPU one: embedding materialises a whole @@ -276,12 +299,12 @@ The app/serving pool is separate and unaffected (5, paired with a matching ### The disk guard Before any bytes are downloaded, each worker checks free space on the filesystem -it is about to write to. Below 2.5 GB it fails **that repo**, not the run: +it is about to write to. Below 0.5 GB it fails **that repo**, not the run: ``` ERROR indexer.job [acme/leviathan]: failed to index acme/leviathan -OSError: insufficient local disk for acme/leviathan: 1904214016 bytes free at /tmp/tmpXXXX, -need 2500000000 (...); lower index_concurrency in config.yaml +OSError: insufficient local disk for acme/leviathan: 104214016 bytes free at /tmp/tmpXXXX, +need 500000000 (...); lower index_concurrency in config.yaml ``` A shortfall fails that repo alone, not the run, and the job exits non-zero. @@ -290,12 +313,12 @@ and re-run — the completed repos are skipped, so the retry is cheap (see §1). **This guard is a pre-flight sanity check, not admission control.** It reserves nothing: each worker calls `shutil.disk_usage` independently before it writes, -so at 5 GB free with 4 workers all four pass their check and all four then +so at 1 GB free with 4 workers all four pass their check and all four then download. It reliably catches the *steady-state* case — disk already low when a repo starts — and turns it into the legible error above. It does **not** bound the *transient* case, where the combined footprint exhausts the disk mid-flight; that still surfaces as an opaque `tarfile` error. Sizing `index_concurrency` to -your actual disk (2.5 GB per worker peak) is the real control. +your actual disk (0.5 GB per worker peak) is the real control. --- diff --git a/indexer/AGENTS.md b/indexer/AGENTS.md index f894434..f052873 100644 --- a/indexer/AGENTS.md +++ b/indexer/AGENTS.md @@ -4,7 +4,7 @@ # indexer ## Purpose -The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` entry point in `pyproject.toml`). It reads the central `config.yaml` from the Databricks workspace, resolves it into a deduped list of GitHub repos (fail-fast on empty or oversized results, before any tarball is fetched or database connection opened), then fans repos out over a bounded thread pool. Per repo it resolves the default branch's immutable HEAD SHA, resolves configured branch globs into a concrete branch list, and — sequentially per branch — downloads the tarball by SHA over plain HTTPS (no git binary), extracts it safely, parses text files, extracts tree-sitter symbols and reference edges (typed call/import sites, Python-only for now), and writes everything in one atomic per-(repo, branch) transaction with content-SHA-deduped storage and a mark-and-sweep of stale branch membership. When semantic search is enabled, files are also chunked and embedded via `app.embed` — outside the transaction — and precomputed vectors are written through a `chunk_writer` seam. After every worker has joined, a post-fan-out checkpoint reconciles desired state — retiring stale branches and purging removed repos — but ONLY on a fully clean run (no failures, conflicts, or truncated branch discovery anywhere); a large repo-purge shrink is withheld as an incident signal rather than applied. The process exits non-zero if any branch fails, if reconciliation itself fails partway, or if a purge was withheld. +The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` entry point in `pyproject.toml`). It reads the central `config.yaml` from the Databricks workspace, resolves it into a deduped list of GitHub repos (fail-fast on empty or oversized results, before any tarball is fetched or database connection opened), then fans repos out over a bounded thread pool. Per repo it resolves the default branch's immutable HEAD SHA, resolves configured branch globs into a concrete branch list, and — sequentially per branch — downloads the tarball by SHA over plain HTTPS (no git binary), streams the text files straight out of that archive in a single in-memory pass (nothing is ever extracted to disk), extracts tree-sitter symbols and reference edges (typed call/import sites, Python-only for now), and writes everything in one atomic per-(repo, branch) transaction with content-SHA-deduped storage and a mark-and-sweep of stale branch membership. When semantic search is enabled, files are also chunked and embedded via `app.embed` — outside the transaction — and precomputed vectors are written through a `chunk_writer` seam. After every worker has joined, a post-fan-out checkpoint reconciles desired state — retiring stale branches and purging removed repos — but ONLY on a fully clean run (no failures, conflicts, or truncated branch discovery anywhere); a large repo-purge shrink is withheld as an incident signal rather than applied. The process exits non-zero if any branch fails, if reconciliation itself fails partway, or if a purge was withheld. ## Key Files | File | Description | @@ -12,12 +12,13 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` | `__init__.py` | Empty package marker. | | `branches.py` | Resolves a connection's branch globs against a repo's real branch list via `fnmatchcase`, returning a typed `BranchResolution` (`branches`, `complete`, `dropped`, `cap`). Default branch always included and always first; empty globs means default-branch-only with no GitHub branches API call, always `complete=True`. `SOFT_BRANCH_CAP = 20` truncates (loud warning naming the cap and dropped branches, not a failure; sets `complete=False`, blocking reconciliation for that repo — #61); no override flag — fix the config. | | `chunk_store.py` | `write_chunks`: delete-and-reinsert one file's rows in the `chunks` table (no natural key). Takes a live `Connection`; never opens an engine, never calls the embedder — vectors arrive precomputed. No `repo_id` parameter: `chunks` is scoped by `file_id` only. | -| `fetch.py` | All GitHub HTTP: paginated org/user repo enumeration (`RepoMeta`), branch listing, `resolve_ref`/`resolve_branch_head` (branch name -> immutable SHA), streamed tarball download capped at `MAX_TARBALL_BYTES` (500 MB), safe extraction (`filter="data"`, bomb check) capped at `MAX_EXTRACTED_BYTES` (2 GB), `assert_disk_headroom` (2.5 GB per worker, both caps alive at once). `RateLimitError` is deliberately narrow: 429 always; 403 only with `Retry-After` or `X-RateLimit-Remaining: 0` — other 403s are permission failures. | +| `fetch.py` | All GitHub HTTP: paginated org/user repo enumeration (`RepoMeta`), branch listing, `resolve_ref`/`resolve_branch_head` (branch name -> immutable SHA), streamed tarball download capped at `MAX_TARBALL_BYTES` (500 MB), `assert_disk_headroom` (0.5 GB per worker — the compressed tarball is the only artifact written to disk). No extraction: `MAX_EXTRACTED_BYTES` moved to `ingest.py` with #106 and is now a work cap, not a disk cap. `RateLimitError` is deliberately narrow: 429 always; 403 only with `Retry-After` or `X-RateLimit-Remaining: 0` — other 403s are permission failures. | +| `ingest.py` | `iter_tar_source_files` (#106): the job's only file source. ONE forward pass over the downloaded tarball (`tarfile.open(mode="r:*")` driven by `tf.next()`; no `getmembers()`, no `extractall()`, nothing written to disk), yielding the same `ParsedFile` stream `parse.iter_source_files` would have yielded from an extracted tree — same filter chain, reusing `parse._looks_binary` by import so the binary-sniff rule cannot fork. `size` is `member.size` (the archive-declared length, which exceeds `len(content)` after NUL stripping); `lang` is the RELATIVE path's lowered suffix. Yields in ARCHIVE order (deterministic per commit; nothing downstream depends on the sequence — see the module docstring's D5 note). `tf.members.clear()` is the FIRST statement of the loop body, since the `TarFile` now outlives `engine.connect()`. Replaces `filter="data"` with explicit, probe-verified rules: raw-`..`-component and absolute member names raise, absolute/escaping link targets raise for hardlinks and symlinks alike, special files and benign hardlinks skip with a WARNING, duplicate names are first-wins, and a zero-member archive raises rather than silently stamping the branch as indexed. `MAX_EXTRACTED_BYTES` (2 GB) is enforced incrementally over pre-filter member sizes. | | `hashing.py` | `content_sha`: canonical SHA-256 hex of content (`None` -> empty string). Single source of truth for `files.content_sha`; must stay byte-identical to the `0003` migration's SQL backfill forever or cross-branch dedup silently breaks (`tests/integration/test_content_sha_parity.py` is the gate). | -| `job.py` | Entry point and orchestration: `main()`/`run()`, `ThreadPoolExecutor` sized by `effective_workers` (unit of work = one repo, all branches sequential), batched `repo_branches` stamp read, per-branch skip-if-unchanged, per-branch `BranchOutcome` classification (`indexed`/`skipped`/`conflict`/`failed`), semantic precompute (`_precompute_chunk_writer`), `ContextVar`-based `[repo]` log attribution, and per-phase timing instrumentation (one `phase timing repo@branch: total=… resolve=… download=… extract=… parse=… embed=… db=… sweep=… other=…` INFO line per INDEXED branch, emitted from inside the worker so `[repo]` resolves; fixed unconditional field set; `_timed_items` charges lazy item production to `parse` so it is not fused into `db`; `resolve=`/`list=` on the per-repo `finished` line carry the repo-scoped costs that belong to no branch's total). Returns 1 if any branch failed; conflicts self-heal and do not fail the run. **Reconciliation checkpoint**: after every worker joins, `_decide_reconciliation` gates on zero failures/conflicts, every repo accounted for, and every repo's `discovery_complete`; on a pass, `_reconcile` reuses the pre-fan-out stamp snapshot to compute each repo's retired branches and the run's full desired repo set, then calls the injected `reconcile_retired_fn`/`reconcile_removed_fn` seams (default: the real `indexer.store` primitives) on one post-fan-out connection — retired branches first, repo purge second. `MAX_PURGE_SHRINK_FRACTION = 0.5` (hardcoded, no config knob) withholds ONLY the purge, as a logged incident signal, if it would remove a strict majority of currently stored repos; retired-branch cleanup on survivors still applies. `ReconcileProgress.committed_any` drives an honest "partially reconciled" vs "left stale" failure message — never the raw exception, only its phase/repo/type name. Never called from worker-thread code (`_index_one*`) — a source-level tripwire test enforces it. | +| `job.py` | Entry point and orchestration: `main()`/`run()`, `ThreadPoolExecutor` sized by `effective_workers` (unit of work = one repo, all branches sequential), batched `repo_branches` stamp read, per-branch skip-if-unchanged, per-branch `BranchOutcome` classification (`indexed`/`skipped`/`conflict`/`failed`), semantic precompute (`_precompute_chunk_writer`), `ContextVar`-based `[repo]` log attribution, and per-phase timing instrumentation (one `phase timing repo@branch: total=… resolve=… download=… parse=… embed=… db=… sweep=… other=…` INFO line per INDEXED branch, emitted from inside the worker so `[repo]` resolves; fixed unconditional field set; `_timed_items` charges lazy item production to `parse` so it is not fused into `db`; `resolve=`/`list=` on the per-repo `finished` line carry the repo-scoped costs that belong to no branch's total). Returns 1 if any branch failed; conflicts self-heal and do not fail the run. **Reconciliation checkpoint**: after every worker joins, `_decide_reconciliation` gates on zero failures/conflicts, every repo accounted for, and every repo's `discovery_complete`; on a pass, `_reconcile` reuses the pre-fan-out stamp snapshot to compute each repo's retired branches and the run's full desired repo set, then calls the injected `reconcile_retired_fn`/`reconcile_removed_fn` seams (default: the real `indexer.store` primitives) on one post-fan-out connection — retired branches first, repo purge second. `MAX_PURGE_SHRINK_FRACTION = 0.5` (hardcoded, no config knob) withholds ONLY the purge, as a logged incident signal, if it would remove a strict majority of currently stored repos; retired-branch cleanup on survivors still applies. `ReconcileProgress.committed_any` drives an honest "partially reconciled" vs "left stale" failure message — never the raw exception, only its phase/repo/type name. Never called from worker-thread code (`_index_one*`) — a source-level tripwire test enforces it. | | `languages.py` | Single source of truth shared by `parse.py` and `symbols.py`: `EXT_TO_LANG` (8 extensions -> tree-sitter language names), `SYMBOL_KINDS` (node type -> symbol kind per language), `EDGE_NODE_KINDS` (node type -> reference-edge kind per language, Python-only until #85), `MAX_FILE_BYTES` (1 MB), `SEMANTIC_CHUNK_MAX_CHARS` (2000, ~4 chars/token), and the frozen dataclasses `ParsedFile`, `Chunk`, `ExtractedSymbol`, `ExtractedEdge`, `FileExtraction`, `IndexCounts` (now carries `edges`). | -| `parse.py` | `iter_source_files`: walk an extracted tree, yield every text file (unknown extensions kept with `lang=None`, since grep runs over all files); skips `.git/`, symlinks, files > `MAX_FILE_BYTES` (stat before read), NUL-sniffed binaries, and UTF-8 decode failures; strips surviving NULs (Postgres `text` rejects them). `iter_chunks`: deterministic line-aligned chunking, no overlap, no mid-line splits, 1-based inclusive line ranges. | -| `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (2.5 GB peak per worker), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. Deliberately import-light: pydantic + PyYAML + stdlib only. | +| `parse.py` | `iter_source_files`: **no production caller since #106** — retained deliberately as the extraction ORACLE the streaming path (`ingest.py`) is pinned against by `tests/unit/test_ingest_parity.py`, and undeletable anyway without editing a `SEMANTICS_PATHS` module and forcing an `INDEX_SEMANTICS_VERSION` bump. Walk an extracted tree, yield every text file (unknown extensions kept with `lang=None`, since grep runs over all files); skips `.git/`, symlinks, files > `MAX_FILE_BYTES` (stat before read), NUL-sniffed binaries, and UTF-8 decode failures; strips surviving NULs (Postgres `text` rejects them). `iter_chunks`: deterministic line-aligned chunking, no overlap, no mid-line splits, 1-based inclusive line ranges. | +| `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (0.5 GB peak per worker since #106), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. Deliberately import-light: pydantic + PyYAML + stdlib only. | | `resolve.py` | `resolve_repos`: enumerate org/user selectors, apply `ExcludeRules` to enumerated repos only (explicit `repos` entries always win, unfiltered), dedup case-insensitively keeping first-seen spelling, union each connection's `branches:` globs into per-repo `RepoEntry.branch_globs`, then fail fast: `EmptyConfigError` on zero repos (indexing nothing must not exit 0), `RepoCeilingError` above `MAX_REPOS` (500, overridable via `--max_repos`). Also matches `RepoConfig.semantic_max_chunks_per_repo` case-insensitively onto each resolved repo's `RepoEntry.semantic_max_chunks` (`None` if unmatched); an override key matching no resolved repo logs a WARNING (typo guard). Enumerator params are a test seam only, not provider dispatch. | | `store.py` | `index_repo`: the single atomic unit of work for one (repo, branch) inside `with conn.begin():` — repos upsert, `repo_branches` CAS-baseline read, per-file array-union upsert on `(repo_id, path, content_sha)` with branch-membership union, delete-and-reinsert `symbols` AND `reference_edges` (both keyed only by `file_id`, no natural key), optional `chunk_writer` call, membership sweep (strip this branch from unseen rows, delete rows with empty `branches`; skipped with a WARNING on an empty seen-set), then CAS stamp (raises `StaleIndexError` on mismatch, rolling everything back). `items` is `Iterable[tuple[ParsedFile, FileExtraction]]`; the `reference_edges` delete runs unconditionally, even when a file's `FileExtraction.edges` is empty, so stale rows never survive a re-index. The sweep call site is wall-clocked into `indexer.timing`'s ambient `PhaseTimer` (`record("sweep", …)`) — never into `IndexCounts` (frozen, compared by value) and never through a new `index_repo` parameter (that signature is an injected seam); `store.py` emits **no** log record for it. Pure DML, no TEMP tables (job role has no TEMP privilege on Lakebase). `reconcile_retired_branches` and `reconcile_removed_repos`: the desired-state storage primitives `job.py`'s reconciliation checkpoint calls — each its own `conn.begin()`, repo-scoped membership subtraction (never delete-by-path) and an exact-match repo purge respectively; neither decides WHAT is retired/desired, only applies a caller-supplied set (see `job.py`'s row). | | `symbols.py` | `extract_file` (#84): one tree-sitter parse, one full-tree walk, emitting both `symbols` and `reference_edges` candidates (`FileExtraction`) in the same pass — nested definitions captured, named nodes only for symbols, kinds from `SYMBOL_KINDS`/`EDGE_NODE_KINDS`, 1-based lines. Stack entries carry the innermost named enclosing definition so edges attribute to it in O(1) with no second walk. `extract_symbols` is a thin wrapper (`extract_file(pf).symbols`) kept for the existing call sites. Python-only edge helpers (`_python_call_target`, `_python_import_edges`) are the seam #85 generalizes to other languages. Parser cache is per-thread (`threading.local`) as insurance against a future GIL-releasing `parse()`. | @@ -39,7 +40,7 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - **Import discipline**: `repo_config.py` stays import-light (no httpx/SQLAlchemy/SDK at module level); `normalize_repo` is imported from `repo_config`, never from `job` (import cycle); `app.embed`'s databricks-sdk dependency is only imported when semantic is on. ### Testing Requirements -- `make test` (`pytest -m "unit or observability"`, no external deps): `tests/unit/test_branches.py`, `test_chunk_store.py`, `test_chunking.py`, `test_edges.py`, `test_fetch.py`, `test_job.py`, `test_job_redaction.py`, `test_languages.py`, `test_parse.py`, `test_repo_config.py`, `test_resolve.py`, `test_symbols.py`, `test_store_chunk_writer.py`, `test_semantics_version_tripwire.py`, `test_timing.py`. +- `make test` (`pytest -m "unit or observability"`, no external deps): `tests/unit/test_branches.py`, `test_chunk_store.py`, `test_chunking.py`, `test_edges.py`, `test_fetch.py`, `test_ingest.py`, `test_ingest_parity.py`, `test_job.py`, `test_job_redaction.py`, `test_languages.py`, `test_parse.py`, `test_repo_config.py`, `test_resolve.py`, `test_symbols.py`, `test_store_chunk_writer.py`, `test_semantics_version_tripwire.py`, `test_timing.py`. - `make test-integration` (needs Postgres): `tests/integration/test_store.py`, `test_store_chunk_writer.py`, `test_content_sha_parity.py` (the hard gate on `hashing.py` vs the migration backfill), `test_reconcile.py` (the `reconcile_retired_branches`/`reconcile_removed_repos` storage primitives directly), and `test_job_reconcile.py` (#59: the real `run()` end to end — real primitives, real engine, only GitHub HTTP faked). - Schema tests must stay fast: changes to `repo_config.py` must not add heavy imports. @@ -48,7 +49,7 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - `pg_insert(...).on_conflict_do_update(...).returning(...)` with a no-op `SET` on conflict — `DO NOTHING ... RETURNING` returns no row on conflict and would break the id bootstrap. - Delete-and-reinsert for child rows with no natural key (symbols, chunks). - `fnmatchcase`, never `fnmatch`: plain globs must behave identically on every platform. -- Guardrail constants with config-level fixes, not override flags: `SOFT_BRANCH_CAP` (20), `MAX_REPOS` (500, the one exception — `--max_repos`), `MAX_FILE_BYTES`, `MAX_TARBALL_BYTES`, `MAX_EXTRACTED_BYTES`, `MAX_PURGE_SHRINK_FRACTION` (0.5, `job.py` — withholds a corpus-wide repo purge above this shrink fraction; fix is staged config removal, never a knob). +- Guardrail constants with config-level fixes, not override flags: `SOFT_BRANCH_CAP` (20), `MAX_REPOS` (500, the one exception — `--max_repos`), `MAX_FILE_BYTES`, `MAX_TARBALL_BYTES` (`fetch.py`, a disk cap), `MAX_EXTRACTED_BYTES` (`ingest.py` since #106, a WORK cap — nothing is written to disk, so it bounds streamed content, not storage), `MAX_PURGE_SHRINK_FRACTION` (0.5, `job.py` — withholds a corpus-wide repo purge above this shrink fraction; fix is staged config removal, never a knob). - Errors name the fix: exceptions carry the config key or job parameter to change (`exclude.size_mb`, `index_concurrency`, `--max_repos`). - Broad `except Exception` at classification boundaries only (`run()`'s config/resolve handlers, `_index_one_branch`), each with a comment explaining why broad is load-bearing. - Lazy generators through the open transaction for bounded memory (semantic mode is the exception — it materializes the file list because embedding needs all chunk texts up front). @@ -68,6 +69,6 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - `pydantic` + `pyyaml` (config schema and parsing) - `tree-sitter` / `tree-sitter-language-pack` (symbol extraction) - `databricks-sdk` (workspace config read + secrets, injected/lazy — never imported at module level here) -- stdlib: `tarfile` (extraction with `filter="data"`), `concurrent.futures`, `tempfile`, `shutil`, `fnmatch`, `hashlib`, `argparse`, `contextvars` +- stdlib: `tarfile` (single-pass in-memory streaming in `ingest.py`; no `extractall`, so no `filter="data"` — the rejection rules are explicit there instead), `posixpath` (member-name and link-target safety), `concurrent.futures`, `tempfile`, `shutil`, `fnmatch`, `hashlib`, `argparse`, `contextvars` diff --git a/indexer/fetch.py b/indexer/fetch.py index 98ad87d..631542a 100644 --- a/indexer/fetch.py +++ b/indexer/fetch.py @@ -2,7 +2,7 @@ The caller sets the ``Authorization`` header on the :class:`httpx.Client`; this module never reads secrets itself. ``download_tarball`` is always called with the -immutable resolved ``sha`` (not a branch name) so the extracted tree's SHA can +immutable resolved ``sha`` (not a branch name) so the indexed tree's SHA can never drift from the ``head_sha`` stamped into ``files.commit``. """ @@ -10,7 +10,6 @@ import logging import shutil -import tarfile import time from dataclasses import dataclass from pathlib import Path @@ -26,16 +25,18 @@ "X-GitHub-Api-Version": "2022-11-28", } -# Defense-in-depth caps for the untrusted tarball on an ephemeral serverless disk: -# bound the compressed download and the uncompressed extraction independently so a -# gzip bomb or an oversized tracked blob can't exhaust local storage. +# Defense-in-depth cap for the untrusted tarball on an ephemeral serverless disk: +# bound the compressed download so an oversized repo can't exhaust local storage. +# The companion cap on decompressed content lives in `indexer.ingest` as +# MAX_EXTRACTED_BYTES -- since #106 nothing is ever extracted to disk, so that one +# bounds WORK (a gzip bomb streamed through memory), not storage. MAX_TARBALL_BYTES = 500_000_000 -MAX_EXTRACTED_BYTES = 2_000_000_000 -# One worker's worst-case peak. The two caps SUM rather than max(): the compressed -# tarball stays on disk inside the worker's TemporaryDirectory while the extracted -# tree grows beside it, so both are alive simultaneously. -REQUIRED_FREE_BYTES = MAX_TARBALL_BYTES + MAX_EXTRACTED_BYTES +# One worker's worst-case peak, and it is now just the download: the compressed +# tarball is the only artifact written to disk, and it is streamed straight out of +# the worker's TemporaryDirectory into `indexer.ingest.iter_tar_source_files` +# without ever being expanded beside it. +REQUIRED_FREE_BYTES = MAX_TARBALL_BYTES _PER_PAGE = 100 @@ -225,10 +226,13 @@ def assert_disk_headroom(path: Path, *, repo: str) -> None: """Raise ``OSError`` unless ``path``'s filesystem can hold one worker's peak. Called immediately before the download, on the directory actually being - written to, so the measurement is of the right filesystem. The error names - the repo AND the config key to lower, because the alternative -- an opaque - ENOSPC from somewhere inside tarfile -- says nothing about which of N - concurrent workers overcommitted the disk or what to do about it. + written to, so the measurement is of the right filesystem. That peak is the + compressed tarball and nothing else: since #106 the archive is streamed + in-memory by :func:`indexer.ingest.iter_tar_source_files` and is never + extracted. The error names the repo AND the config key to lower, because the + alternative -- an opaque ENOSPC from somewhere inside tarfile -- says nothing + about which of N concurrent workers overcommitted the disk or what to do + about it. The caller is expected to let this fail ONE repo, not the run: with a too-high ``index_concurrency`` the run then degrades to indexing whatever @@ -238,8 +242,8 @@ def assert_disk_headroom(path: Path, *, repo: str) -> None: if free < REQUIRED_FREE_BYTES: raise OSError( f"insufficient local disk for {repo}: {free} bytes free at {path}, need " - f"{REQUIRED_FREE_BYTES} (a {MAX_TARBALL_BYTES}-byte tarball plus a " - f"{MAX_EXTRACTED_BYTES}-byte extraction, both alive at once); " + f"{REQUIRED_FREE_BYTES} (the compressed tarball is the only artifact " + "written to disk; it is never extracted); " "lower index_concurrency in config.yaml" ) @@ -269,33 +273,3 @@ def download_tarball(client: httpx.Client, org: str, repo: str, ref: str, dest: ) fh.write(chunk) return out - - -def extract_tarball(tar_path: Path, dest: Path) -> Path: - """Safely extract ``tar_path`` into ``dest`` and return its single top-level dir. - - ``filter="data"`` (Python 3.12) neutralizes path traversal / absolute paths / - device files. GitHub tarballs contain exactly one top-level ``org-repo-/`` - directory. - """ - dest.mkdir(parents=True, exist_ok=True) - with tarfile.open(tar_path, mode="r:*") as tf: - members = tf.getmembers() - # Reject a decompression bomb before writing any of it to disk. - extracted = sum(m.size for m in members if m.isreg()) - if extracted > MAX_EXTRACTED_BYTES: - raise ValueError( - f"tarball extracts to {extracted} bytes, exceeding {MAX_EXTRACTED_BYTES}" - ) - tf.extractall(dest, filter="data") - - top_level = { - member.name.split("/", 1)[0] - for member in members - if member.name and not member.name.startswith("/") - } - if len(top_level) != 1: - raise ValueError( - f"expected exactly one top-level dir in tarball, found {sorted(top_level)}" - ) - return dest / next(iter(top_level)) diff --git a/indexer/ingest.py b/indexer/ingest.py new file mode 100644 index 0000000..e630cf3 --- /dev/null +++ b/indexer/ingest.py @@ -0,0 +1,365 @@ +"""Stream a downloaded repo tarball once and yield the text files worth indexing. + +This is the streaming counterpart of :func:`indexer.parse.iter_source_files`, and +the only file source the indexing job uses (issue #106). It exists for two +reasons: + +* **One decompression, not two.** ``extract_tarball`` used to call + ``tf.getmembers()`` (a full gzip decompression, to build the member list) and + then ``tf.extractall()`` (a second full decompression, plus a write of every + byte to disk). This module walks the stream exactly once with ``tf.next()``: + there is no ``getmembers()`` and no ``extractall()`` anywhere in ``indexer/``. +* **No extracted tree.** Nothing is written to disk, so one worker's peak local + disk is the compressed tarball alone (``MAX_TARBALL_BYTES``), not the tarball + plus its expansion. See ``indexer.fetch.REQUIRED_FREE_BYTES``. + +**Ordering contract: ARCHIVE order, and the contract is determinism plus set +equality -- never sequence equality.** There is no cheap metadata pass over a +gzip stream (``tarfile`` locates each header by walking the decompressed bytes, +so a names-only pass costs the very decompression this module removes), and +sorting would mean materialising every ``ParsedFile`` in memory, destroying the +bounded-memory property the non-semantic path depends on. A GitHub tarball is +``git archive`` order -- git tree order for a fixed commit -- so the sequence is +deterministic per ``head_sha``. Nothing downstream depends on it: +``indexer.store``'s per-file classification and its sweep are both order-free. + +**``indexer.parse.iter_source_files`` is this module's executable oracle.** It has +no production caller any more, and is deliberately kept: ``tests/unit/ +test_ingest_parity.py`` extracts each fixture tarball to disk, runs +``iter_source_files`` over the tree and this module over the same tarball, and +requires identical ``(path, lang, size, content)`` sets. The expected value is +computed live from ``parse.py`` rather than from a golden fixture, so it cannot +be regenerated to match a drifting implementation. ``parse.py`` is also watched +by ``tests/unit/test_semantics_version_tripwire.py``; retiring it is a deliberate +semantics-version conversation and is out of scope here. + +The filter chain below MUST stay identical to ``parse.py``'s: ``.git/`` skip, +``MAX_FILE_BYTES`` skip, NUL-in-the-first-8-KB binary sniff, UTF-8 decode or +skip, strip surviving NULs. ``_looks_binary`` and ``_BINARY_SNIFF_BYTES`` are +imported from ``parse`` rather than re-implemented -- private on purpose, since a +forked copy of the sniff rule is exactly how the two would silently diverge (the +8 KB window itself rides along inside it; ``_BINARY_SNIFF_BYTES`` is deliberately +NOT imported separately, as an unused import). +Two details are load-bearing for that parity and are *not* self-evident: +``size`` is the archive-declared ``member.size`` (the analogue of +``stat().st_size``), not ``len(content)`` -- it exceeds it after NUL stripping -- +and ``lang`` is derived from the *relative* path's lowered suffix, after the top +directory is stripped. + +Behavioural divergence from the old extract-then-walk path, PROBE-VERIFIED +against ``extract_tarball`` + ``iter_source_files`` on Python 3.12 (several of +these contradict what the ``tarfile`` docs suggest; do not re-derive them from +memory): + +===================================== ================================== ========================== +Member shape Old (extract-to-disk, verified) New (stream) +===================================== ================================== ========================== +benign internal symlink extracted, then skipped skipped -- identical +symlink with an absolute target ``AbsoluteLinkError``, branch ``ValueError`` -- same + fails fail-loud posture +symlink escaping the archive ``LinkOutsideDestinationError`` ``ValueError`` -- same +hard link, benign internal target materialised as a file and skipped + WARNING + INDEXED (git has no hardlinks) +hard link, absolute/escaping target same link errors as symlinks ``ValueError`` -- same +device / FIFO / socket ``SpecialFileError``, branch skipped + WARNING + fails (nothing is written now) +``/etc/passwd`` leading ``/`` stripped, extracted ``ValueError`` + outside the top dir, SILENTLY + IGNORED +``/{TOP}/evil.py`` leading ``/`` stripped -- INDEXED ``ValueError`` +``../evil.txt`` ``OutsideDestinationError`` ``ValueError`` +``{TOP}/../evil.txt`` normalised to ``evil.txt``, ``ValueError`` + SILENTLY IGNORED +``{TOP}/../{TOP}/evil.py`` normalised back inside -- ``ValueError`` + INDEXED +``./{TOP}/a.py`` top dir NOT stripped, yielded as ``a.py`` + ``{TOP}/a.py`` +empty member name ``IsADirectoryError`` ``ValueError`` +regular file named exactly ``{TOP}`` yielded nothing ``ValueError`` +duplicate member name last wins first wins + WARNING +two or more top-level dirs ``ValueError`` ``ValueError`` +zero members / no top-level dir ``ValueError`` ``ValueError`` +===================================== ================================== ========================== + +The four silent rows (two absolute names, two ``..`` names) are latent +path-confusion bugs this module closes: two members were silently dropped and +two were silently INDEXED under a path the archive did not really contain. +Absolute and escaping link targets deliberately keep today's hard failure rather +than relaxing to a skip -- nothing is written to disk any more so a skip would be +safe, but changing which repos index is out of scope here. + +Validation (the bomb cap, the top-level-dir check, path safety, any ``tarfile`` +corruption error, and the trailing-stream drain) now happens as the stream is +consumed. On the non-semantic path that is inside ``index_repo``'s open +transaction, so a malformed archive surfaces as a rolled-back transaction rather +than a pre-connection failure. The branch-level outcome is unchanged +(``status="failed"``). +""" + +from __future__ import annotations + +import logging +import posixpath +import tarfile +from collections.abc import Iterator +from pathlib import Path + +from indexer.languages import EXT_TO_LANG, MAX_FILE_BYTES, ParsedFile + +# Private on purpose: the binary-sniff rule must be un-forkable, so it is +# imported from the oracle module rather than re-implemented here. Precedent: +# tests/unit/test_parse.py already imports _BINARY_SNIFF_BYTES. +from indexer.parse import _looks_binary + +logger = logging.getLogger("indexer.ingest") + +# Upper bound on the uncompressed content one branch may stream out of its +# tarball. Moved here from indexer.fetch by #106, and re-scoped with it: nothing +# is written to disk any more, so this is a WORK cap (a decompression-bomb +# guard), not a disk cap. Enforced incrementally, two ways: +# +# * every regular member's declared size, accumulated BEFORE the +# .git/size/binary filters -- exactly what the old up-front +# `sum(m.size for m in members if m.isreg())` counted, and still the tighter +# check for the case that matters (one huge tracked blob); +# * `member.offset`, the tar stream's cumulative decompressed position, which +# bounds members of ANY type. Regular-file accounting alone would let an +# archive of a million directory or link headers decompress unbounded -- +# they carry no data, so `streamed` never moves. +MAX_EXTRACTED_BYTES = 2_000_000_000 + + +def _normalise_member_name(name: str) -> str: + """Return ``name`` with a leading ``./`` stripped, rejecting unsafe shapes. + + The ONLY normalisation performed is stripping a leading ``./`` (GNU tar + writes those; the old path left them on and consequently failed to strip the + top-level directory at all). Everything else raises. + + ``..`` is rejected as a **raw path component**, before any normalisation, and + ``posixpath.normpath`` is deliberately NOT used here -- normalising first + hides two of the three traversal shapes the D7 table requires to raise:: + + raw name normpath() normpath-then-check + ../evil.txt ../evil.txt raise + {TOP}/../evil.txt evil.txt ALLOW <- wrong + {TOP}/../{TOP}/evil.py {TOP}/evil.py ALLOW <- wrong, and it + then passes the top-dir + check and is indexed + + For a member *name* the question is "does it contain ``..`` at all", not + "where does it resolve to"; ``normpath`` belongs only in the link-target + check, where the second question is the real one. The empty name needs its + own test because ``normpath("")`` returns ``"."``. + """ + if name.startswith("./"): + name = name[2:] + if not name: + raise ValueError("tarball contains a member with an empty name") + if posixpath.isabs(name): + raise ValueError(f"tarball member has an absolute name: {name!r}") + if ".." in name.split("/"): + raise ValueError(f"tarball member name contains a '..' component: {name!r}") + return name + + +def _assert_link_target_is_contained(name: str, linkname: str) -> None: + """Raise ``ValueError`` if ``linkname`` is absolute or escapes the archive. + + Pure ``posixpath`` string arithmetic, NOT a transliteration of CPython's + ``filter="data"`` check: that one computes + ``os.path.realpath(os.path.join(dest, os.path.dirname(name), linkname))``, + and there is no destination directory in a stream -- substituting a + placeholder would resolve against the real filesystem and the process cwd, + which is a silently wrong check rather than a strict one. + + Diverges from ``filter="data"`` only for links chained through other + symlinks, which ``git archive`` cannot produce. + """ + if posixpath.isabs(linkname): + raise ValueError(f"tarball member {name!r} links to an absolute target: {linkname!r}") + resolved = posixpath.normpath(posixpath.join(posixpath.dirname(name), linkname)) + if resolved == ".." or resolved.startswith("../"): + raise ValueError(f"tarball member {name!r} links outside the archive: {linkname!r}") + + +def iter_tar_source_files(tar_path: Path) -> Iterator[ParsedFile]: + """Yield a :class:`ParsedFile` per indexable text file in ``tar_path``. + + A drop-in replacement for ``indexer.parse.iter_source_files(root)``: same + return type, same element semantics, one argument. ONE forward pass over the + archive -- no ``getmembers()``, no ``extractall()``, nothing written to disk. + + ``path`` is repo-relative (the single top-level ``org-repo-/`` + directory GitHub tarballs carry is stripped). Skips directories, links, + special files, ``.git/`` contents, members declared larger than + ``MAX_FILE_BYTES``, duplicate paths, binaries (NUL sniff) and UTF-8 decode + failures. Raises ``ValueError`` for an unsafe member name or link target, for + an archive with anything other than exactly one top-level directory, and when + the streamed content exceeds :data:`MAX_EXTRACTED_BYTES`. + + Peak memory is ``MAX_FILE_BYTES`` plus the seen-path set (measured with + ``tracemalloc`` at ~120 bytes per path -- the set retains the path STRING, + not just a slot -- so ~12 MB at 100k files), held for as long as the consumer + holds the generator, which on the non-semantic path is the whole open + transaction. + + ``mode="r:gz"``, not ``"r:*"``: GitHub only ever serves gzip, and pinning the + codec removes bzip2/LZMA from the attack surface entirely. Their compression + ratios are far higher than DEFLATE's, so an ``r:*`` reader would accept a + much smaller upload for the same decompressed volume. + """ + tf = tarfile.open(tar_path, mode="r:gz") + try: + top_dir: str | None = None + streamed = 0 + seen: set[str] = set() + + while (member := tf.next()) is not None: + # FIRST statement of the body, not the last: TarFile.next() appends + # every TarInfo to tf.members, and unlike extract_tarball (whose + # member list died before engine.connect()) this TarFile stays open + # for the whole transaction. At the END of the body every `continue` + # below would skip it -- which is most members in a real repo -- and + # the list would grow unbounded. Safe because next() only appends + # and extractfile() reads member.offset_data, never the list; the + # `while ... tf.next()` loop is what makes it safe (TarFile.__iter__ + # indexes into self.members and would be corrupted by this). + # + # The ignore is typeshed's gap, not a runtime one: `TarFile.members` + # is a plain list assigned in `TarFile.__init__` and appended to by + # `next()`, but the stub only declares `getmembers()` (which would + # read the WHOLE archive -- exactly what AC1 removes). + tf.members.clear() # type: ignore[attr-defined] + + name = _normalise_member_name(member.name) + + component = name.split("/", 1)[0] + if top_dir is None: + top_dir = component + elif component != top_dir: + raise ValueError( + "expected exactly one top-level dir in tarball, found " + f"{sorted({top_dir, component})}" + ) + + # Before the isreg() gate, because links are isreg()-false. The old + # path ran filter="data"'s link check on `islnk() or issym()` alike, + # so hardlinks are in scope for the same fail-loud rule. + if member.islnk() or member.issym(): + _assert_link_target_is_contained(name, member.linkname) + + # BEFORE the isreg() gate, so it bounds members of EVERY type -- + # including the ones that carry no data and so never move `streamed` + # below. `member.offset` is this header's position in the DECOMPRESSED + # tar stream and is monotonic across a forward pass, so it is the + # honest measure of how much has actually been decompressed. Without + # it an archive of a million directory or link headers decompresses + # unbounded past a cap that only ever sums regular-file sizes. + if member.offset > MAX_EXTRACTED_BYTES: + raise ValueError( + f"tarball stream reaches {member.offset} decompressed bytes at member " + f"{name!r}, exceeding {MAX_EXTRACTED_BYTES}" + ) + + if not member.isreg(): + if member.islnk(): + # Divergence: extractall() materialised these as real files + # and they WERE indexed. git archives contain no hardlinks. + logger.warning("skipping hard link member %s in %s", name, tar_path.name) + elif not (member.isdir() or member.issym()): + # Divergence: extractall() raised SpecialFileError and failed + # the branch. Nothing is written to disk now, so the reason + # to fail loudly is gone. + logger.warning("skipping special file member %s in %s", name, tar_path.name) + continue + + # Counted BEFORE the filters below, exactly like the old up-front + # `sum(m.size for m in members if m.isreg())`: same threshold, same + # population, only the timing differs. Still the tighter check for a + # single oversized member, whose declared size is caught from its + # header before any of its data is read. + streamed += member.size + if streamed > MAX_EXTRACTED_BYTES: + raise ValueError( + f"tarball streams to {streamed} bytes of content, " + f"exceeding {MAX_EXTRACTED_BYTES}" + ) + + rel_path = name[len(top_dir) + 1 :] if name != top_dir else "" + if not rel_path: + raise ValueError(f"tarball member is the top-level dir itself: {name!r}") + + if ".git" in rel_path.split("/"): + continue + + # From the header alone -- the analogue of parse.py's stat() before + # read_bytes(), so an oversized blob's data is never read at all. + if member.size > MAX_FILE_BYTES: + continue + + if rel_path in seen: + # extractall() was last-wins; a single forward pass cannot look + # ahead, so this is first-wins. git archives have no duplicates. + logger.warning("skipping duplicate member %s in %s", name, tar_path.name) + continue + seen.add(rel_path) + + fh = tf.extractfile(member) + if fh is None: + # Unreachable: the isreg() gate above already excludes every + # member for which extractfile() returns None. This narrows + # tarfile's declared `IO[bytes] | None` for mypy; it is not a + # real branch. + continue + with fh: + raw = fh.read() + + if _looks_binary(raw): + continue + try: + content = raw.decode("utf-8") + except UnicodeDecodeError: + continue + + # Postgres `text` rejects NUL (0x00), which is legal UTF-8 and so can + # survive the 8 KB sniff window (issue #37). Stripped before the + # yield so stored content and its content_sha stay consistent -- + # `size` stays the archive-declared length and may exceed + # len(content), matching parse.py's on-disk size. + content = content.replace("\x00", "") + + yield ParsedFile( + path=rel_path, + lang=EXT_TO_LANG.get(Path(rel_path).suffix.lower()), + size=member.size, + content=content, + ) + + # A truncated archive is NOT reliably a hard error on the read path: + # `TarFile.next()` only re-raises a truncated/invalid header when + # `self.offset == 0`, so a cut mid-stream can simply return None and end + # the loop early. That yields a PARTIAL file set, which is worse than a + # zero-file one -- the missing files are absent from `index_repo`'s seen + # set, so the membership sweep deletes them from the corpus. Draining the + # remaining bytes forces gzip's CRC32/ISIZE trailer check, which a + # truncated or corrupt stream cannot pass. On a well-formed archive this + # reads only the trailing block padding. + fileobj = tf.fileobj + if fileobj is not None: + while fileobj.read(65536): + pass + + if top_dir is None: + # A zero-member (empty) archive. This raise is the one guard against a + # bug that is invisible in production: without it the branch indexes + # zero files, index_repo skips its sweep on the empty-seen-set guard, + # and _stamp_repo_branch(seen_any=False) writes (head_sha, unchanged + # version) -- which for a branch already stamped at the current + # INDEX_SEMANTICS_VERSION matches the skip seam exactly, silently + # marking it current at a HEAD it never read. + raise ValueError("expected exactly one top-level dir in tarball, found none") + finally: + # A consumer that abandons this generator (gen.close(), or an exception + # out of index_repo's transaction) still releases the file handle. + tf.close() diff --git a/indexer/job.py b/indexer/job.py index 704209f..b4ccb3f 100644 --- a/indexer/job.py +++ b/indexer/job.py @@ -10,7 +10,9 @@ Orchestrates, per resolved repo: resolve the default branch's HEAD -> resolve the repo's concrete branch list from the config globs -> for each branch, SEQUENTIALLY: resolve its HEAD SHA -> download the tarball by that immutable SHA --> extract -> parse text files -> extract symbols -> atomic upsert + mark-and-sweep +-> stream text files straight out of that archive (:func:`indexer.ingest. +iter_tar_source_files`, one pass, nothing extracted to disk) -> extract symbols +-> atomic upsert + mark-and-sweep via :func:`indexer.store.index_repo`. Branches within one repo are sequential (never concurrent) -- that is the invariant that keeps ``store.py``'s per-branch sweep sound without an advisory lock. Each BRANCH is isolated: one branch's @@ -98,7 +100,7 @@ and the aggregate run-completion WARNING this module emits for it. Every INDEXED branch also emits one ``phase timing`` line accounting for its whole -wall clock -- resolve / download / extract / parse / embed / db / sweep, plus an +wall clock -- resolve / download / parse / embed / db / sweep, plus an ``other`` residual so no time is silently unattributed. The fields are fixed and unconditional (a phase that did not run prints ``0.00s``) so the line stays greppable whether or not semantic indexing is on. Skipped, failed, and conflicted @@ -143,14 +145,14 @@ REQUIRED_FREE_BYTES, assert_disk_headroom, download_tarball, - extract_tarball, list_branches, resolve_branch_head, resolve_ref, ) from indexer.hashing import content_sha +from indexer.ingest import iter_tar_source_files from indexer.languages import Chunk, FileExtraction, IndexCounts, ParsedFile -from indexer.parse import iter_chunks, iter_source_files +from indexer.parse import iter_chunks from indexer.repo_config import RepoConfig, effective_workers, load_config, normalize_repo from indexer.resolve import MAX_REPOS, RepoEntry, resolve_repos from indexer.store import ( @@ -1180,10 +1182,6 @@ def _index_one_branch( tar_path = download_tarball(http_client, org, repo, head_sha, tmp_path) timer.add("download", timer.clock() - t0) - t0 = timer.clock() - root = extract_tarball(tar_path, tmp_path / "extracted") - timer.add("extract", timer.clock() - t0) - chunk_writer: ChunkWriter | None = None precompute_failed = False if cfg.semantic_enabled and embed_fn is not None: @@ -1192,13 +1190,15 @@ def _index_one_branch( # index_repo's open transaction. # # This walk is charged to `parse` explicitly: it is the same - # rglob + stat + read + decode work that _timed_items charges on + # stream + decompress + decode work that _timed_items charges on # the non-semantic path, and leaving it unwrapped would dump the # entire file-walk cost of the PRODUCTION (semantic-on) path into # `other`, which is exactly the number this instrumentation - # exists to route work by. + # exists to route work by. Since #106 it also carries the gzip + # decompression that used to be reported as its own `extract` + # phase -- there is no separate extraction step any more. t0 = timer.clock() - files = list(iter_source_files(root)) + files = list(iter_tar_source_files(tar_path)) timer.add("parse", timer.clock() - t0) # In a `finally`, unlike every other phase wrap: the degrade path @@ -1207,7 +1207,7 @@ def _index_one_branch( # advisory shas_fn read is charged here too, deliberately NOT as # its own timed phase: it exists solely to decide what this block # embeds, and #103's `phase timing` line is pinned exhaustive - # (nine fixed fields, tests/unit/test_job.py) -- adding a tenth + # (eight fixed fields, tests/unit/test_job.py) -- adding a ninth # field is out of this change's scope. t0 = timer.clock() try: @@ -1266,7 +1266,7 @@ def _index_one_branch( items = ((pf, extract_file(pf)) for pf in files) else: # Lazy generator: files stream through the open transaction (bounded memory). - items = ((pf, extract_file(pf)) for pf in iter_source_files(root)) + items = ((pf, extract_file(pf)) for pf in iter_tar_source_files(tar_path)) # Parse time is INSIDE the db window (items are produced lazily as # index_repo consumes them), so `db` subtracts only the parse accrued @@ -1290,8 +1290,9 @@ def _index_one_branch( ) db_wall = timer.clock() - t0 - # Emitted HERE -- after the TemporaryDirectory teardown (an rm -rf of a - # possibly multi-GB extracted tree, which `other` must include) and from + # Emitted HERE -- after the TemporaryDirectory teardown (since #106 an + # rm -rf of one compressed tarball rather than a multi-GB extracted + # tree, but still time `other` must include) and from # inside the worker, where _index_one's _repo_ctx still resolves the # [%(repo)s] field. The drain loop on the main thread would render `[-]`. # @@ -1315,7 +1316,6 @@ def _index_one_branch( total - timer.total("resolve") - timer.total("download") - - timer.total("extract") - timer.total("parse") - timer.total("embed") - db @@ -1323,16 +1323,20 @@ def _index_one_branch( ) # One format string, no branches: a phase that did not run prints 0.00s # rather than vanishing, so the line stays greppable and field-stable - # whether or not semantic indexing is on. + # whether or not semantic indexing is on. That is a promise about one + # build's runs, NOT that the field set is immutable across releases: + # #106 removed `extract=` because the phase ceased to exist (the + # archive is streamed, never extracted), not because it could read + # zero. A field that can never be non-zero is dead weight that sends + # an operator hunting for a phase that is not there. logger.info( - "phase timing %s@%s: total=%.2fs resolve=%.2fs download=%.2fs extract=%.2fs " + "phase timing %s@%s: total=%.2fs resolve=%.2fs download=%.2fs " "parse=%.2fs embed=%.2fs db=%.2fs sweep=%.2fs other=%.2fs", name, branch, total, timer.total("resolve"), timer.total("download"), - timer.total("extract"), timer.total("parse"), timer.total("embed"), db, diff --git a/indexer/repo_config.py b/indexer/repo_config.py index 8327bf3..4d39a9f 100644 --- a/indexer/repo_config.py +++ b/indexer/repo_config.py @@ -245,14 +245,17 @@ class RepoConfig(BaseModel): ``index_concurrency`` is how many repos the indexing job works on at once. **The default of 4 is a disk bound, not a CPU one.** Each in-flight worker - holds ``MAX_TARBALL_BYTES`` (500 MB) *and* ``MAX_EXTRACTED_BYTES`` (2 GB) - alive simultaneously -- the downloaded tarball stays inside the worker's - ``TemporaryDirectory`` while the extraction runs beside it, so peak usage is - 2.5 GB per worker: 10 GB at the default 4, 20 GB at the ceiling of 8. + holds ``MAX_TARBALL_BYTES`` (500 MB) inside its ``TemporaryDirectory``, and + that is the whole of its on-disk footprint: since #106 the archive is + streamed once in memory by ``indexer.ingest.iter_tar_source_files`` and never + extracted, so ``MAX_EXTRACTED_BYTES`` (2 GB, now in ``indexer.ingest``) caps + WORK rather than storage and does not add to this. Peak usage is 0.5 GB per + worker: 2 GB at the default 4, 4 GB at the ceiling of 8. (#106 lowered those + numbers 5x but deliberately left the default alone; #109 re-derives it.) **Returns at the ceiling are sublinear.** Symbol extraction does not parallelise (measured at 0.95x on 4 threads), so Amdahl's law caps the - speedup well below 8x while the disk cost stays a hard linear 20 GB. Raise + speedup well below 8x while the disk cost stays a hard linear 4 GB. Raise it only knowing that trade. When semantic indexing is on, the effective worker count is clamped to 2 by diff --git a/indexer/timing.py b/indexer/timing.py index 4a7f477..45d549c 100644 --- a/indexer/timing.py +++ b/indexer/timing.py @@ -1,7 +1,7 @@ """Ambient per-phase wall-clock accounting for one branch's indexing pipeline. -``indexer.job`` measures most of a branch's phases (resolve / download / extract / -parse / embed / db) inline, but ``sweep`` runs deep inside +``indexer.job`` measures most of a branch's phases (resolve / download / parse / +embed / db) inline, but ``sweep`` runs deep inside :func:`indexer.store.index_repo`, behind the injected ``index_fn`` seam and a frozen ``IndexCounts`` return type. Threading a timer through that seam would force every existing ``index_fn`` fake to grow a parameter, turning a log-only diff --git a/tests/integration/test_job_ingest_delta.py b/tests/integration/test_job_ingest_delta.py new file mode 100644 index 0000000..27a7510 --- /dev/null +++ b/tests/integration/test_job_ingest_delta.py @@ -0,0 +1,147 @@ +"""Streaming ingestion does not rewrite an already-indexed corpus (#106 AC3 + #104). + +This is the end-to-end proof that #106's byte-identical-corpus claim survives +contact with #104's file-level delta gate. AC3 forbids an +``INDEX_SEMANTICS_VERSION`` bump, so the delta gate is **open** on the first run +after #106 deploys -- which means every file whose ``(path, content_sha)`` still +matches issues no statement at all. If ``ingest.py`` diverged from ``parse.py`` in +any content-affecting way, that first run would instead reclassify the whole +corpus as changed and rewrite it. + +The ``lang``/``size`` half of that guarantee is NOT observable here and must not +be assumed from a green run: the gate keys on ``(path, content_sha)`` alone, so a +``lang``/``size`` divergence would ALSO show up as all-unchanged, silently and +permanently. That dimension is pinned by the oracle assertions in +``tests/unit/test_ingest_parity.py``. + +Clones ``tests/integration/test_store_delta.py``'s throwaway-schema fixture idiom +(own copy, per ``tests/integration/AGENTS.md``'s no-conftest convention), and for +the same reason builds no ``chunks`` table and creates no ``lakebase_*`` +extension -- that is what lets it run against a plain local Postgres. +""" + +from __future__ import annotations + +import logging +import tarfile +from collections.abc import Iterator +from pathlib import Path + +import pytest +from sqlalchemy import Connection, text + +from app.db.client import create_db_engine +from app.db.models import Base +from indexer.ingest import iter_tar_source_files +from indexer.languages import IndexCounts +from indexer.parse import iter_source_files +from indexer.store import index_repo +from indexer.symbols import extract_file +from tests.unit.test_ingest import TOP, _dir, _Entry, _reg, _write + +SCHEMA = "test_job_ingest_delta" +REPO = "acme/widgets" + +# Deliberately mixed: mapped and unmapped extensions, a nested directory, and two +# members BOTH paths must skip -- so an "all unchanged" result also proves the two +# implementations agree on what is NOT in the corpus. +_FIXTURE: list[_Entry] = [ + _dir(TOP), + _reg(f"{TOP}/main.py", b"def f():\n return 1\n"), + _reg(f"{TOP}/pkg/util.py", b"class C:\n pass\n"), + _reg(f"{TOP}/app.js", b"function g() { return 2; }\n"), + _reg(f"{TOP}/README.md", b"# hello\n"), + _reg(f"{TOP}/.git/config", b"[core]\n"), + _reg(f"{TOP}/logo.png", b"\x89PNG\x00\x00binary"), +] + + +@pytest.fixture +def conn() -> Iterator[Connection]: + engine = create_db_engine() + connection = engine.connect() + try: + connection.execute(text(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE")) + connection.execute(text(f"CREATE SCHEMA {SCHEMA}")) + connection.execute(text(f"SET search_path TO {SCHEMA}, public")) + connection.commit() + + Base.metadata.create_all(bind=connection) + connection.commit() + + yield connection + finally: + connection.rollback() + connection.execute(text(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE")) + connection.commit() + connection.close() + engine.dispose() + + +def _delta_lines(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + r.getMessage() + for r in caplog.records + if r.name == "indexer.store" and "delta write set" in r.getMessage() + ] + + +@pytest.mark.integration +def test_reindexing_an_already_indexed_branch_is_all_unchanged( + conn: Connection, tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """Index through the OLD path, re-index the SAME tarball through the new one. + + Leg 1 extracts to disk and walks with ``parse.iter_source_files``, exactly as + ``extract_tarball`` + the pre-#106 job did. Leg 2 streams the same archive + with ``iter_tar_source_files`` at a new ``head_sha``. Every file must classify + as *unchanged*: zero writes, and every ``files.id``/``symbols.id`` preserved + (a delete-reinsert would renumber the serials, so identical ids are the + precise proof). + """ + tar_path = _write(tmp_path, _FIXTURE) + + dest = tmp_path / "extracted" + dest.mkdir() + with tarfile.open(tar_path, mode="r:*") as tf: + tf.extractall(dest, filter="data") + old_items = [(pf, extract_file(pf)) for pf in iter_source_files(dest / TOP)] + + first = index_repo( + conn, + name=REPO, + branch="main", + is_default=True, + head_sha="sha_first", + items=old_items, + ) + assert first.files == len(old_items) == 4 # .git/config and the PNG are skipped + assert first.symbols > 0 + + files_before = dict(conn.execute(text("SELECT path, id FROM files ORDER BY path")).all()) + symbols_before = sorted(conn.execute(text("SELECT id FROM symbols")).scalars().all()) + conn.rollback() + + with caplog.at_level(logging.INFO, logger="indexer.store"): + second = index_repo( + conn, + name=REPO, + branch="main", + is_default=True, + head_sha="sha_second", + items=((pf, extract_file(pf)) for pf in iter_tar_source_files(tar_path)), + ) + + assert second == IndexCounts(files=first.files, symbols=0, swept=0, edges=0) + + files_after = dict(conn.execute(text("SELECT path, id FROM files ORDER BY path")).all()) + symbols_after = sorted(conn.execute(text("SELECT id FROM symbols")).scalars().all()) + assert files_after == files_before + assert symbols_after == symbols_before + + lines = _delta_lines(caplog) + assert len(lines) == 1 + assert ( + f"delta write set 0/{first.files} files " + f"(unchanged={first.files} membership=0, semantics gate open)" in lines[0] + ) diff --git a/tests/unit/test_fetch.py b/tests/unit/test_fetch.py index ea43c4d..58a31c2 100644 --- a/tests/unit/test_fetch.py +++ b/tests/unit/test_fetch.py @@ -17,7 +17,6 @@ RepoMeta, assert_disk_headroom, download_tarball, - extract_tarball, list_branches, list_org_repos, list_user_repos, @@ -104,46 +103,10 @@ def test_download_tarball_rejects_oversized( download_tarball(client, ORG, REPO, SHA, tmp_path) -@pytest.mark.unit -def test_extract_tarball_rejects_decompression_bomb( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - # Cap below the summed member size -> extraction is refused before writing. - monkeypatch.setattr(fetch, "MAX_EXTRACTED_BYTES", 4) - tar_path = tmp_path / "source.tar.gz" - tar_path.write_bytes(CLEAN_TARBALL) - dest = tmp_path / "extracted" - with pytest.raises(ValueError, match="exceeding"): - extract_tarball(tar_path, dest) - assert not any(dest.iterdir()) if dest.exists() else True - - -@pytest.mark.unit -def test_extract_tarball_yields_top_level_dir(tmp_path: Path) -> None: - tar_path = tmp_path / "source.tar.gz" - tar_path.write_bytes(CLEAN_TARBALL) - root = extract_tarball(tar_path, tmp_path / "extracted") - assert root.name == TOP_DIR - assert (root / "README.md").read_text() == "# hello\n" - assert (root / "src" / "main.py").exists() - - -@pytest.mark.unit -def test_extract_tarball_neutralizes_path_traversal(tmp_path: Path) -> None: - malicious = _make_tarball( - { - f"{TOP_DIR}/ok.py": b"x = 1\n", - "../evil.txt": b"pwned\n", - } - ) - tar_path = tmp_path / "evil.tar.gz" - tar_path.write_bytes(malicious) - dest = tmp_path / "extracted" - - # The data filter blocks the escaping member; nothing is written outside dest. - with pytest.raises(tarfile.OutsideDestinationError): - extract_tarball(tar_path, dest) - assert not (dest.parent / "evil.txt").exists() +# Extraction moved out of this module entirely (#106): the tarball is streamed +# once, in memory, by `indexer.ingest.iter_tar_source_files`. The bomb cap, the +# top-level-dir contract and the path-traversal rejections that used to be tested +# here now live in tests/unit/test_ingest.py, against that function. # --- Enumeration ------------------------------------------------------------- @@ -499,14 +462,16 @@ def test_size_cap_error_names_overage_and_config_key( @pytest.mark.unit -def test_required_free_bytes_sums_both_caps() -> None: - """Both caps are alive at once, so they SUM. - - The tarball stays on disk inside the worker's TemporaryDirectory while the - extracted tree grows beside it; a max() here would under-reserve by 500 MB - per worker and silently reintroduce the failure the guard exists to prevent. +def test_required_free_bytes_is_the_tarball_cap_alone() -> None: + """The compressed tarball is the only artifact on disk (#106). + + It used to be the tarball PLUS its extraction, both alive at once inside the + worker's TemporaryDirectory, so the two caps summed. Nothing is extracted any + more -- `indexer.ingest.iter_tar_source_files` streams the archive in memory + -- so reserving `MAX_EXTRACTED_BYTES` on top would over-reserve 2 GB per + worker and cap `index_concurrency` on disk the job never touches. """ - assert fetch.REQUIRED_FREE_BYTES == fetch.MAX_TARBALL_BYTES + fetch.MAX_EXTRACTED_BYTES + assert fetch.REQUIRED_FREE_BYTES == fetch.MAX_TARBALL_BYTES @pytest.mark.unit @@ -527,7 +492,10 @@ def test_assert_disk_headroom_error_names_the_repo_and_the_config_key( """The whole point of the guard is diagnosability. An opaque ENOSPC from inside tarfile says neither which of N concurrent - workers overcommitted the disk nor what to change, so both are asserted. + workers overcommitted the disk nor what to change, so both are asserted. The + message also has to say WHAT the reserved number is, and since #106 that is + the compressed tarball alone -- an operator reading "plus a 2 GB extraction" + would size `index_concurrency` against a footprint that no longer exists. """ monkeypatch.setattr( fetch.shutil, "disk_usage", lambda _p: _Usage(free=fetch.REQUIRED_FREE_BYTES - 1) @@ -539,6 +507,8 @@ def test_assert_disk_headroom_error_names_the_repo_and_the_config_key( assert f"{ORG}/{REPO}" in message assert "index_concurrency" in message assert str(fetch.REQUIRED_FREE_BYTES) in message + assert "never extracted" in message + assert "extraction, both alive at once" not in message class _Usage: diff --git a/tests/unit/test_ingest.py b/tests/unit/test_ingest.py new file mode 100644 index 0000000..bf0a97d --- /dev/null +++ b/tests/unit/test_ingest.py @@ -0,0 +1,678 @@ +"""Unit tests for indexer.ingest.iter_tar_source_files against in-memory tarballs. + +The behavioural-divergence tests below carry the OLD (extract-to-disk) behaviour in +their docstrings. Those are probe-verified against ``indexer.fetch.extract_tarball`` ++ ``indexer.parse.iter_source_files`` on Python 3.12, not inferred from the +``tarfile`` documentation, which misleads on several of them -- four rows that +sound like they must have raised in fact silently dropped or silently INDEXED a +member. Do not "correct" them from memory. + +Whole-corpus equivalence with ``indexer.parse.iter_source_files`` lives in +``test_ingest_parity.py``; this file pins the edges that parity cannot reach +(members ``filter="data"`` rejects outright). +""" + +from __future__ import annotations + +import io +import logging +import tarfile +from pathlib import Path +from typing import Any + +import pytest + +import indexer.ingest as ingest +from indexer.ingest import iter_tar_source_files +from indexer.languages import MAX_FILE_BYTES + +ORG = "acme" +REPO = "widgets" +SHA = "abc1234def5678" +TOP = f"{ORG}-{REPO}-{SHA[:7]}" + +_Entry = tuple[tarfile.TarInfo, bytes | None] + + +# --- fixture builders (mirrors tests/unit/test_fetch.py's, plus link/dir/special) --- + + +def _reg(name: str, data: bytes) -> _Entry: + info = tarfile.TarInfo(name) + info.type = tarfile.REGTYPE + info.size = len(data) + return info, data + + +def _dir(name: str) -> _Entry: + info = tarfile.TarInfo(name) + info.type = tarfile.DIRTYPE + return info, None + + +def _sym(name: str, target: str) -> _Entry: + info = tarfile.TarInfo(name) + info.type = tarfile.SYMTYPE + info.linkname = target + return info, None + + +def _lnk(name: str, target: str) -> _Entry: + info = tarfile.TarInfo(name) + info.type = tarfile.LNKTYPE + info.linkname = target + return info, None + + +def _fifo(name: str) -> _Entry: + info = tarfile.TarInfo(name) + info.type = tarfile.FIFOTYPE + return info, None + + +def _make_tarball(entries: list[_Entry]) -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + for info, data in entries: + tf.addfile(info, io.BytesIO(data) if data is not None else None) + return buf.getvalue() + + +def _write(tmp_path: Path, entries: list[_Entry], name: str = "source.tar.gz") -> Path: + out = tmp_path / name + out.write_bytes(_make_tarball(entries)) + return out + + +_CLEAN = [ + _dir(TOP), + _reg(f"{TOP}/README.md", b"# hello\n"), + _reg(f"{TOP}/src/main.py", b"def f():\n return 1\n"), +] + + +class _CountingBytesIO(io.BytesIO): + """A seekable in-memory file that records how many bytes were pulled out of it.""" + + def __init__(self, data: bytes) -> None: + super().__init__(data) + self.bytes_read = 0 + + def read(self, size: int | None = -1, /) -> bytes: + chunk = super().read(size) + self.bytes_read += len(chunk) + return chunk + + def read1(self, size: int = -1, /) -> bytes: + chunk = super().read1(size) + self.bytes_read += len(chunk) + return chunk + + +def _capture_tarfiles(monkeypatch: pytest.MonkeyPatch) -> list[tarfile.TarFile]: + """Record every ``TarFile`` ``iter_tar_source_files`` opens. + + ``tf`` is local to the generator, so this monkeypatch on + ``indexer.ingest.tarfile.open`` is the only handle a test can get on it. + """ + opened: list[tarfile.TarFile] = [] + real_open = tarfile.open + + def _spy(*args: Any, **kwargs: Any) -> tarfile.TarFile: + tf = real_open(*args, **kwargs) + opened.append(tf) + return tf + + monkeypatch.setattr(ingest.tarfile, "open", _spy) + return opened + + +# --- AC1 / AC2: one decompression, nothing on disk -------------------------- + + +@pytest.mark.unit +def test_decompresses_the_archive_exactly_once( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """AC1. Measures DECOMPRESSED BYTES PULLED, not ``getmembers()`` calls. + + A spy on ``getmembers`` would pass for an implementation that opens the + archive twice and streams it twice -- exactly the double decompression this + change exists to remove. ``ingest`` calls ``tarfile.open`` with a path, so + there is no file object to wrap: patch ``open`` itself and re-dispatch onto a + counting ``fileobj``. + + The 1.05x threshold is safe rather than sloppy: ``mode="r:gz"`` dispatches + straight to ``gzopen`` with no codec probing (so no failed-probe re-read), and + ``next()`` never seeks backwards, so no ``DecompressReader.rewind()`` occurs. + The end-of-stream drain reads only the trailing block padding. A second pass + would double the count outright. + """ + tar_path = _write(tmp_path, _CLEAN) + compressed = tar_path.read_bytes() + reader = _CountingBytesIO(compressed) + opens: list[tuple[Any, ...]] = [] + real_open = tarfile.open + + def _spy(*args: Any, **kwargs: Any) -> tarfile.TarFile: + opens.append(args) + return real_open(fileobj=reader, mode="r:gz") + + monkeypatch.setattr(ingest.tarfile, "open", _spy) + + files = list(iter_tar_source_files(tar_path)) + + assert [pf.path for pf in files] == ["README.md", "src/main.py"] + assert len(opens) == 1, "the archive was opened more than once" + assert reader.bytes_read <= len(compressed) * 1.05, ( + f"read {reader.bytes_read} bytes from a {len(compressed)}-byte archive; " + "that is a second pass" + ) + + +# --- happy path ------------------------------------------------------------- + + +@pytest.mark.unit +def test_yields_expected_files_from_a_clean_tarball(tmp_path: Path) -> None: + """Top dir stripped, extensions mapped, unknown extensions kept with lang=None.""" + tar_path = _write( + tmp_path, + [ + _dir(TOP), + _reg(f"{TOP}/README.md", b"# hello\n"), + _reg(f"{TOP}/src/main.py", b"def f():\n return 1\n"), + _reg(f"{TOP}/Makefile", b"all:\n"), + ], + ) + by_path = {pf.path: pf for pf in iter_tar_source_files(tar_path)} + + assert set(by_path) == {"README.md", "src/main.py", "Makefile"} + assert by_path["src/main.py"].lang == "python" + assert by_path["src/main.py"].content == "def f():\n return 1\n" + assert by_path["src/main.py"].size == len(b"def f():\n return 1\n") + assert by_path["README.md"].lang is None + assert by_path["Makefile"].lang is None + + +@pytest.mark.unit +def test_ordering_is_deterministic(tmp_path: Path) -> None: + """D5: archive order, and archive order is fixed for a fixed commit. + + The contract downstream is determinism plus set-equality, never a sorted + sequence -- sorting would need either a metadata pass (the decompression AC1 + removes) or the whole corpus in memory. + """ + entries = [_dir(TOP)] + [_reg(f"{TOP}/{n}.py", b"x = 1\n") for n in ("z", "a", "m", "b")] + tar_path = _write(tmp_path, entries) + + first = [pf.path for pf in iter_tar_source_files(tar_path)] + second = [pf.path for pf in iter_tar_source_files(tar_path)] + + assert first == second + assert first == ["z.py", "a.py", "m.py", "b.py"] + + +@pytest.mark.unit +def test_size_is_the_archive_declared_length_not_len_content(tmp_path: Path) -> None: + """``size`` is ``member.size`` -- parse.py's ``stat().st_size`` analogue. + + A NUL past the 8 KB sniff window is legal UTF-8, decodes cleanly, and is then + stripped (Postgres ``text`` rejects it), so the stored content is SHORTER than + the file. Using ``len(content)`` here would diverge from ``parse.py`` on every + such file -- and #104's delta gate keys only on ``(path, content_sha)``, so + that divergence would classify as *unchanged* and never be corrected. + """ + data = b"a" * 9000 + b"\x00" + b"b" + tar_path = _write(tmp_path, [_dir(TOP), _reg(f"{TOP}/nul.py", data)]) + + (pf,) = list(iter_tar_source_files(tar_path)) + assert pf.size == len(data) + assert pf.content == "a" * 9000 + "b" + assert pf.size > len(pf.content) + + +# --- the caps --------------------------------------------------------------- + + +@pytest.mark.unit +def test_oversized_member_is_skipped_without_reading_its_data( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """D9: the size check is on the header, so a huge blob's data is never read. + + ``member.size`` replaces ``entry.stat().st_size`` for exactly this reason -- + peak memory stays bounded by ``MAX_FILE_BYTES``, not by the largest member. + """ + extracted: list[str] = [] + real_extractfile = tarfile.TarFile.extractfile + + def _spy(self: tarfile.TarFile, member: Any) -> Any: + extracted.append(member.name) + return real_extractfile(self, member) + + monkeypatch.setattr(tarfile.TarFile, "extractfile", _spy) + + tar_path = _write( + tmp_path, + [ + _dir(TOP), + _reg(f"{TOP}/huge.py", b"x" * (MAX_FILE_BYTES + 1)), + _reg(f"{TOP}/small.py", b"x = 1\n"), + ], + ) + assert [pf.path for pf in iter_tar_source_files(tar_path)] == ["small.py"] + assert extracted == [f"{TOP}/small.py"] + + +@pytest.mark.unit +def test_incremental_cap_raises_at_the_same_threshold( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """D9: same threshold as the old up-front ``sum(m.size ...)``, different timing. + + The single member sits at offset 0 so the coarser ``member.offset`` guard + cannot fire first -- this test is specifically about the regular-file size + accumulator, which stays the tighter of the two for one oversized blob. The + message carries the actual magnitude, not just the cap, so an operator can + tell "repo slightly over" from "gzip bomb". + """ + monkeypatch.setattr(ingest, "MAX_EXTRACTED_BYTES", 4) + tar_path = _write(tmp_path, [_reg(f"{TOP}/big.py", b"0123456789")]) + + with pytest.raises(ValueError, match="streams to 10 bytes of content, exceeding 4"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_cap_counts_filtered_members_too(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """D9: pre-filter accounting, matching ``sum(m.size for m in members if m.isreg())``. + + The old bomb check ran before (and independently of) the ``.git``/size/binary + filters. An archive whose bulk is ``.git/`` objects is still a bomb, so the + counter is incremented before any of them. The blob is the FIRST member, at + offset 0, so this pins the size accumulator rather than the coarser + ``member.offset`` guard. + """ + monkeypatch.setattr(ingest, "MAX_EXTRACTED_BYTES", 4) + tar_path = _write(tmp_path, [_reg(f"{TOP}/.git/objects/pack/blob", b"0123456789")]) + + with pytest.raises(ValueError, match="streams to 10 bytes of content, exceeding 4"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_offset_cap_bounds_members_that_carry_no_data( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The regular-file accumulator alone does NOT bound decompression. + + Directory, link and special members carry no data, so ``streamed`` never + moves for them -- an archive of a million such headers would decompress + unbounded past a cap that only ever sums ``member.size`` for ``isreg()`` + members. ``member.offset`` (the tar stream's cumulative decompressed + position) is checked BEFORE the ``isreg()`` gate for exactly this reason. + """ + monkeypatch.setattr(ingest, "MAX_EXTRACTED_BYTES", 1024) + entries: list[_Entry] = [_dir(TOP)] + [_dir(f"{TOP}/d{n}") for n in range(8)] + tar_path = _write(tmp_path, entries) + + with pytest.raises(ValueError, match="decompressed bytes at member"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +@pytest.mark.parametrize("keep", [0.3, 0.6, 0.9]) +def test_truncated_archive_raises_instead_of_yielding_a_partial_set( + tmp_path: Path, keep: float +) -> None: + """A cut-short download must fail the branch, never index a prefix of it. + + ``TarFile.next()`` re-raises a truncated header only when ``self.offset == + 0``; past the first member it can simply return ``None`` and end the loop. + That is worse than yielding nothing: the files beyond the cut are absent from + ``index_repo``'s seen set, so the membership sweep DELETES them from the + corpus, and the branch is then stamped as current at a HEAD it only partly + read. Draining the remaining bytes after the loop forces gzip's CRC32/ISIZE + trailer check, which a truncated stream cannot pass. + """ + entries: list[_Entry] = [_dir(TOP)] + [ + _reg(f"{TOP}/f{n}.py", (f"def f{n}():\n return {n}\n" * 40).encode()) for n in range(60) + ] + tar_path = _write(tmp_path, entries) + assert len(list(iter_tar_source_files(tar_path))) == 60 # intact: all present + + cut = tmp_path / "cut.tar.gz" + cut.write_bytes(tar_path.read_bytes()[: int(tar_path.stat().st_size * keep)]) + + with pytest.raises((EOFError, tarfile.ReadError, ValueError)): + list(iter_tar_source_files(cut)) + + +@pytest.mark.unit +def test_members_list_does_not_grow_across_the_stream( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """D6/D9: ``tf.members.clear()`` is the FIRST statement of the loop body. + + ``TarFile.next()`` appends every ``TarInfo`` to ``self.members``, and unlike + ``extract_tarball`` (whose list died before ``engine.connect()``) this + ``TarFile`` lives for the whole open transaction. At the END of the loop body + the clear would be skipped by every ``continue`` -- which is most members in a + real repo -- and the list would grow unbounded. + """ + entries: list[_Entry] = [_dir(TOP)] + for n in range(30): + entries.append(_reg(f"{TOP}/.git/obj{n}", b"skipped\n")) + entries.append(_reg(f"{TOP}/f{n}.py", b"x = 1\n")) + # Built BEFORE the spy is installed -- `_write` opens a TarFile of its own. + tar_path = _write(tmp_path, entries) + opened = _capture_tarfiles(monkeypatch) + + for _pf in iter_tar_source_files(tar_path): + assert len(opened) == 1 + assert len(opened[0].members) <= 1 + + assert len(opened[0].members) <= 1 + + +# --- D7: member names ------------------------------------------------------- + + +@pytest.mark.unit +def test_absolute_member_name_raises_where_today_it_was_silently_dropped( + tmp_path: Path, +) -> None: + """PROBE-VERIFIED old behaviour: NO error. + + ``filter="data"`` stripped the leading ``/`` and extracted ``/etc/passwd`` to + ``dest/etc/passwd``; ``fetch.py``'s top-level scan then excluded the member + (it started with ``/``), so the file landed outside ``root`` and was SILENTLY + IGNORED. Now it fails the branch, naming the member. + """ + tar_path = _write( + tmp_path, [_dir(TOP), _reg(f"{TOP}/ok.py", b"x = 1\n"), _reg("/etc/passwd", b"root\n")] + ) + + with pytest.raises(ValueError, match="absolute name"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_absolute_member_name_under_the_top_dir_raises(tmp_path: Path) -> None: + """PROBE-VERIFIED old behaviour: NO error, and ``evil.py`` was INDEXED. + + The POSIX filter stripped the leading ``/`` from ``/{TOP}/evil.py``, which + landed it *inside* ``root`` under a path the archive never actually declared. + This is the latent path-confusion bug the raise closes. + """ + tar_path = _write(tmp_path, [_dir(TOP), _reg(f"/{TOP}/evil.py", b"pwned = 1\n")]) + + with pytest.raises(ValueError, match="absolute name"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_parent_traversal_member_name_raises(tmp_path: Path) -> None: + """Old behaviour: ``tarfile.OutsideDestinationError``, branch failed. Now ValueError.""" + tar_path = _write(tmp_path, [_dir(TOP), _reg("../evil.txt", b"pwned\n")]) + + with pytest.raises(ValueError, match=r"'\.\.' component"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_parent_traversal_inside_the_destination_raises(tmp_path: Path) -> None: + """PROBE-VERIFIED old behaviour: NO error, SILENTLY IGNORED. + + ``{TOP}/../evil.txt`` normalises to ``evil.txt``, which is inside ``dest`` -- + so ``filter="data"`` passed it -- but outside ``root``, so the walk never saw + it. The check here is on the RAW components precisely so this row raises: + normalising first would return ``evil.txt`` and allow it. + """ + tar_path = _write(tmp_path, [_dir(TOP), _reg(f"{TOP}/../evil.txt", b"pwned\n")]) + + with pytest.raises(ValueError, match=r"'\.\.' component"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_parent_traversal_back_under_the_top_dir_raises(tmp_path: Path) -> None: + """PROBE-VERIFIED old behaviour: NO error, and ``evil.py`` was INDEXED. + + ``{TOP}/../{TOP}/evil.py`` normalises back to ``{TOP}/evil.py``. This is the + dangerous row: a ``posixpath.normpath``-then-check implementation would + ALLOW it, and it would then pass the top-dir check and be indexed as + ``evil.py`` -- silently reproducing the exact bug this change claims to close. + """ + tar_path = _write(tmp_path, [_dir(TOP), _reg(f"{TOP}/../{TOP}/evil.py", b"pwned = 1\n")]) + + with pytest.raises(ValueError, match=r"'\.\.' component"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_empty_member_name_raises(tmp_path: Path) -> None: + """Old behaviour: ``IsADirectoryError`` out of ``extractall``. Now a named ValueError. + + ``posixpath.normpath("")`` returns ``"."``, which is why the empty name needs + its own check rather than falling out of the traversal one. + """ + tar_path = _write(tmp_path, [_reg("", b"")]) + + with pytest.raises(ValueError, match="empty name"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_member_named_exactly_the_top_dir_raises(tmp_path: Path) -> None: + """A REGULAR file named exactly ``{TOP}`` leaves an empty relative path. + + Old behaviour: ``rglob`` on a non-directory yielded nothing, so the branch + indexed zero files. Streaming would otherwise yield ``path=""``. + """ + tar_path = _write(tmp_path, [_reg(TOP, b"not a directory\n")]) + + with pytest.raises(ValueError, match="top-level dir itself"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_two_top_level_dirs_raise(tmp_path: Path) -> None: + """GitHub tarballs carry exactly one ``org-repo-/``; anything else is not one.""" + tar_path = _write( + tmp_path, + [_dir(TOP), _reg(f"{TOP}/a.py", b"x = 1\n"), _reg("other/b.py", b"y = 2\n")], + ) + + with pytest.raises(ValueError, match="exactly one top-level dir"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_empty_archive_raises(tmp_path: Path) -> None: + """The one bug that would be INVISIBLE in production, so it gets an explicit raise. + + Without it a truncated or empty archive yields zero files, ``index_repo`` + skips its sweep on the empty-seen-set guard, and ``_stamp_repo_branch( + seen_any=False)`` writes ``(head_sha, unchanged version)``. For a branch + already stamped at the current ``INDEX_SEMANTICS_VERSION`` -- the production + steady state -- that matches the skip seam exactly, silently marking the + branch current at a HEAD whose corpus was never read, until HEAD moves again. + #104's ``seen_any`` guard narrows this to that case; it does NOT remove it. + """ + tar_path = _write(tmp_path, []) + + with pytest.raises(ValueError, match="found none"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_leading_dot_slash_is_normalised(tmp_path: Path) -> None: + """PROBE-VERIFIED old behaviour: the top dir was NOT stripped. + + ``./{TOP}/a.py`` extracted to ``dest/{TOP}/a.py``, but ``fetch.py``'s + top-level scan saw ``.`` as the single top-level name and returned ``dest`` + itself as ``root`` -- so every path came out prefixed with ``{TOP}/``. + """ + tar_path = _write(tmp_path, [_dir(f"./{TOP}"), _reg(f"./{TOP}/a.py", b"x = 1\n")]) + + assert [pf.path for pf in iter_tar_source_files(tar_path)] == ["a.py"] + + +# --- D7: link targets (fail-loud posture preserved) ------------------------- + + +@pytest.mark.unit +def test_symlink_with_an_absolute_target_raises(tmp_path: Path) -> None: + """Preserves today's ``tarfile.AbsoluteLinkError``: the branch still fails. + + Nothing is written to disk any more, so a skip would be safe -- but changing + which repos index is out of #106's scope. Relaxing it is a separate issue. + """ + tar_path = _write(tmp_path, [_dir(TOP), _sym(f"{TOP}/link.py", "/etc/passwd")]) + + with pytest.raises(ValueError, match="absolute target"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_symlink_escaping_the_archive_raises(tmp_path: Path) -> None: + """Preserves today's ``tarfile.LinkOutsideDestinationError``. + + The check is pure ``posixpath`` string arithmetic. A transliteration of + ``filter="data"``'s ``os.path.realpath`` would resolve against the real + filesystem and the process cwd -- silently wrong, not merely strict. + """ + tar_path = _write(tmp_path, [_dir(TOP), _sym(f"{TOP}/link.py", "../../../etc/passwd")]) + + with pytest.raises(ValueError, match="links outside the archive"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_hardlink_with_an_absolute_target_raises(tmp_path: Path) -> None: + """Hardlinks are in ``filter="data"``'s link check too (it tests ``islnk() or issym()``). + + Probe-verified: ``LNKTYPE {TOP}/b.py -> /etc/passwd`` raises + ``AbsoluteLinkError`` today, so the fail-loud rule applies to hardlinks + identically -- a blanket "hardlink -> skip" would turn a branch failure into + a silent skip. + """ + tar_path = _write(tmp_path, [_dir(TOP), _lnk(f"{TOP}/b.py", "/etc/passwd")]) + + with pytest.raises(ValueError, match="absolute target"): + list(iter_tar_source_files(tar_path)) + + +@pytest.mark.unit +def test_hardlink_escaping_the_archive_raises(tmp_path: Path) -> None: + """Probe-verified: ``LNKTYPE -> ../../../etc/passwd`` raises today too.""" + tar_path = _write(tmp_path, [_dir(TOP), _lnk(f"{TOP}/b.py", "../../../etc/passwd")]) + + with pytest.raises(ValueError, match="links outside the archive"): + list(iter_tar_source_files(tar_path)) + + +# --- D7: documented divergences --------------------------------------------- + + +@pytest.mark.unit +def test_benign_internal_symlink_is_skipped(tmp_path: Path) -> None: + """The no-divergence case: extracted then skipped by ``is_symlink()`` before, skipped now.""" + tar_path = _write( + tmp_path, + [ + _dir(TOP), + _reg(f"{TOP}/real.py", b"x = 1\n"), + _sym(f"{TOP}/link.py", "real.py"), + ], + ) + + assert [pf.path for pf in iter_tar_source_files(tar_path)] == ["real.py"] + + +@pytest.mark.unit +def test_benign_hardlink_member_is_skipped_with_a_warning( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """DOCUMENTED DIVERGENCE, probe-verified: today the hardlink WAS indexed. + + ``extractall`` materialised it as a real regular file, so the walk saw two + files (``['a.py', 'b.py']``). A stream cannot materialise it, so it is skipped + -- loudly, because it is a real (if unreachable) corpus difference. ``git + archive`` never emits hardlinks. + """ + tar_path = _write( + tmp_path, + [_dir(TOP), _reg(f"{TOP}/a.py", b"x = 1\n"), _lnk(f"{TOP}/b.py", f"{TOP}/a.py")], + ) + + with caplog.at_level(logging.WARNING, logger="indexer.ingest"): + assert [pf.path for pf in iter_tar_source_files(tar_path)] == ["a.py"] + + assert any("hard link" in r.getMessage() for r in caplog.records) + + +@pytest.mark.unit +def test_special_file_member_is_skipped_with_a_warning( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """DOCUMENTED DIVERGENCE: ``extractall`` raised ``SpecialFileError`` and failed the branch. + + Nothing is written to disk now, so the reason to fail loudly is gone -- but + the member is still worth a WARNING, since a repo containing one is unusual. + """ + tar_path = _write(tmp_path, [_dir(TOP), _reg(f"{TOP}/a.py", b"x = 1\n"), _fifo(f"{TOP}/pipe")]) + + with caplog.at_level(logging.WARNING, logger="indexer.ingest"): + assert [pf.path for pf in iter_tar_source_files(tar_path)] == ["a.py"] + + assert any("special file" in r.getMessage() for r in caplog.records) + + +@pytest.mark.unit +def test_duplicate_member_name_keeps_the_first_and_warns( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """DOCUMENTED DIVERGENCE: ``extractall`` was last-wins; a forward pass is first-wins. + + A single stream cannot look ahead to discover that a later member shadows an + earlier one, and buffering to find out would defeat the bounded-memory + property. ``git archive`` emits no duplicate names. + """ + tar_path = _write( + tmp_path, + [_dir(TOP), _reg(f"{TOP}/a.py", b"first = 1\n"), _reg(f"{TOP}/a.py", b"second = 2\n")], + ) + + with caplog.at_level(logging.WARNING, logger="indexer.ingest"): + files = list(iter_tar_source_files(tar_path)) + + assert [pf.content for pf in files] == ["first = 1\n"] + assert any("duplicate member" in r.getMessage() for r in caplog.records) + + +# --- resource discipline ---------------------------------------------------- + + +@pytest.mark.unit +def test_tarfile_handle_is_closed_when_the_consumer_abandons_the_generator( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """D6: the ``try/finally`` around the loop, not a ``with`` at the call site. + + The non-semantic path hands this generator straight into ``index_repo``'s open + transaction; an exception there abandons it mid-stream, and the handle must + still be released. + """ + # Built BEFORE the spy is installed -- `_write` opens a TarFile of its own. + tar_path = _write(tmp_path, _CLEAN) + opened = _capture_tarfiles(monkeypatch) + + gen = iter_tar_source_files(tar_path) + assert next(gen).path == "README.md" + assert opened[0].closed is False + + gen.close() + assert opened[0].closed is True diff --git a/tests/unit/test_ingest_parity.py b/tests/unit/test_ingest_parity.py new file mode 100644 index 0000000..7339e3e --- /dev/null +++ b/tests/unit/test_ingest_parity.py @@ -0,0 +1,145 @@ +"""Oracle parity: ``indexer.ingest`` must agree with ``indexer.parse``, file for file. + +``indexer.parse.iter_source_files`` has no production caller after #106. It is kept +deliberately, as the **executable specification** the streaming path is pinned +against: it cannot be deleted without editing ``indexer/parse.py``, which +``tests/unit/test_semantics_version_tripwire.py`` watches and which would force +the ``INDEX_SEMANTICS_VERSION`` bump #106's AC3 forbids. + +The expected value here is computed LIVE from ``parse.py`` -- each fixture tarball +is extracted to disk exactly the way ``extract_tarball`` used to extract it +(``filter="data"``), walked with ``iter_source_files``, and compared against +``iter_tar_source_files`` over the same tarball. A golden-fixture version of this +test could be regenerated to match a drifting ``ingest.py``; this one cannot. Do +not substitute one. + +**Fixture constraint, load-bearing:** every member here must be one +``filter="data"`` survives. The oracle side raises on special files and on +absolute or escaping links (probe-verified), so including them would make this +test red by construction rather than by defect. Those member shapes are pinned in +``test_ingest.py`` instead. +""" + +from __future__ import annotations + +import tarfile +from pathlib import Path + +import pytest + +from indexer.hashing import content_sha +from indexer.ingest import iter_tar_source_files +from indexer.languages import EXT_TO_LANG, MAX_FILE_BYTES, ParsedFile +from indexer.parse import iter_source_files +from tests.unit.test_ingest import TOP, _dir, _Entry, _reg, _sym, _write + +# A NUL past the 8 KB sniff window: legal UTF-8, decodes cleanly, and is stripped +# before the yield -- so `size` (the archive-declared length) exceeds len(content). +_NUL_PAST_SNIFF = b"a" * 9000 + b"\x00" + b"b" + +_RICH_FIXTURE: list[_Entry] = [ + _dir(TOP), + _reg(f"{TOP}/README.md", b"# hello\n"), + _reg(f"{TOP}/empty.py", b""), + _reg(f"{TOP}/Makefile", b"all:\n"), + _reg(f"{TOP}/deeply/nested/dir/mod.py", b"CONST = 1\n"), + _reg(f"{TOP}/café.py", b"# non-ascii filename\n"), + _dir(f"{TOP}/emptydir"), + _reg(f"{TOP}/real.py", b"def f():\n return 1\n"), + _sym(f"{TOP}/link.py", "real.py"), + _reg(f"{TOP}/.git/config", b"[core]\n"), + _reg(f"{TOP}/huge.py", b"x" * (MAX_FILE_BYTES + 1)), + _reg(f"{TOP}/binary.png", b"\x89PNG\x00\x00 not text"), + _reg(f"{TOP}/nul_past_sniff.py", _NUL_PAST_SNIFF), + _reg(f"{TOP}/latin1.txt", b"\xff\xfe caf\xe9"), + # Uppercase suffix: both sides must lower-case before the EXT_TO_LANG lookup + # (`parse.py` uses `entry.suffix.lower()`, `ingest.py` the relative path's). + # Without this entry nothing in the suite would notice a dropped `.lower()`. + _reg(f"{TOP}/UPPER.PY", b"UPPER = 1\n"), +] + [_reg(f"{TOP}/sample{ext}", b"// sample\n") for ext in sorted(EXT_TO_LANG)] + +# Only members BOTH sides skip: `.git/`, oversized, NUL-sniffed binary, invalid +# UTF-8, a directory, and a benign internal symlink. +_UNINDEXABLE_FIXTURE: list[_Entry] = [ + _dir(TOP), + _dir(f"{TOP}/sub"), + _reg(f"{TOP}/.git/config", b"[core]\n"), + _reg(f"{TOP}/huge.py", b"x" * (MAX_FILE_BYTES + 1)), + _reg(f"{TOP}/binary.png", b"\x89PNG\x00\x00 not text"), + _reg(f"{TOP}/latin1.txt", b"\xff\xfe caf\xe9"), + _sym(f"{TOP}/link", ".git/config"), +] + + +def _oracle(tar_path: Path, dest: Path) -> list[ParsedFile]: + """Extract exactly as ``extract_tarball`` did, then walk with ``parse.py``.""" + dest.mkdir(parents=True, exist_ok=True) + with tarfile.open(tar_path, mode="r:*") as tf: + tf.extractall(dest, filter="data") + return list(iter_source_files(dest / TOP)) + + +def _key(pf: ParsedFile) -> tuple[str, str | None, int, str, str]: + return (pf.path, pf.lang, pf.size, pf.content, content_sha(pf.content)) + + +@pytest.mark.unit +def test_parity_with_iter_source_files(tmp_path: Path) -> None: + """Set equality on every field the corpus is built from, over a rich fixture. + + Set equality, not sequence equality: the stream yields in archive order (D5), + which is deterministic per commit but is not ``sorted()``. Nothing downstream + of ``index_repo`` depends on the order. + """ + tar_path = _write(tmp_path, _RICH_FIXTURE) + + expected = _oracle(tar_path, tmp_path / "extracted") + actual = list(iter_tar_source_files(tar_path)) + + assert expected, "the fixture must actually index something" + assert {_key(pf) for pf in actual} == {_key(pf) for pf in expected} + + +@pytest.mark.unit +def test_parity_holds_when_no_file_is_indexable(tmp_path: Path) -> None: + """Both sides must agree on the EMPTY result too, not just on a populated one. + + An archive of nothing but skips is the shape most likely to expose a filter + that fires on one side and not the other. + """ + tar_path = _write(tmp_path, _UNINDEXABLE_FIXTURE) + + assert _oracle(tar_path, tmp_path / "extracted") == [] + assert list(iter_tar_source_files(tar_path)) == [] + + +@pytest.mark.unit +def test_lang_and_size_match_parse_exactly_for_every_fixture_file(tmp_path: Path) -> None: + """D11: ``lang`` and ``size``, asserted field-by-field and on purpose. + + The set-equality test above already covers these, so this looks redundant -- + it is not, and the reason must be said out loud. #104's delta gate classifies + a file on ``(path, content_sha)`` ALONE (``indexer/store.py``). If + ``ingest.py`` ever derived ``size`` or ``lang`` differently from ``parse.py``, + every already-stored file would classify as *unchanged*, no ``UPDATE`` would + be issued, and the stored columns would never be corrected: content parity + would hold perfectly while ``lang``/``size`` rotted silently and permanently. + ``store.py`` reasons it is protected from that by the tripwire watching the + modules that derive them -- and #106 moved that derivation into ``ingest.py``. + So: ``size`` is ``member.size`` (the archive-declared length, which exceeds + ``len(content)`` after NUL stripping), and ``lang`` comes from the RELATIVE + path's lowered suffix, after the top directory is stripped. + """ + tar_path = _write(tmp_path, _RICH_FIXTURE) + + expected = {pf.path: pf for pf in _oracle(tar_path, tmp_path / "extracted")} + actual = {pf.path: pf for pf in iter_tar_source_files(tar_path)} + + assert set(actual) == set(expected) + for path, pf in actual.items(): + assert pf.lang == expected[path].lang, f"lang diverged for {path}" + assert pf.size == expected[path].size, f"size diverged for {path}" + + # The fixture really does exercise both edges this test exists for. + assert actual["nul_past_sniff.py"].size > len(actual["nul_past_sniff.py"].content) + assert {pf.lang for pf in actual.values()} >= set(EXT_TO_LANG.values()) diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index de21d19..f232286 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -25,6 +25,7 @@ import tarfile import threading import time +from pathlib import Path from typing import Any, NamedTuple import httpx @@ -119,10 +120,15 @@ def __init__( files: dict[str, bytes] | None = None, branches: dict[str, list[str]] | None = None, branches_fail: set[str] | None = None, + tarball_bytes: bytes | None = None, ) -> None: self.enumerations = enumerations or {} self.missing = missing or set() self.files = files + # Verbatim archive bytes for the /tarball endpoint, bypassing `_tarball`. + # The only way to serve a MALFORMED archive, which is what proves where + # streaming validation now happens (D8). + self.tarball_bytes = tarball_bytes # full_name -> branch names, for the /branches endpoint. A repo not # listed here answers with an empty branch list. self.branches = branches or {} @@ -167,6 +173,8 @@ def __call__(self, request: httpx.Request) -> httpx.Response: names = self.branches.get(f"{org}/{repo}", []) return httpx.Response(200, json=[{"name": n} for n in names]) if parts[3] == "tarball": + if self.tarball_bytes is not None: + return httpx.Response(200, content=self.tarball_bytes) return httpx.Response(200, content=_tarball(f"{org}-{repo}-shashas", self.files)) return httpx.Response(404) @@ -2613,6 +2621,146 @@ def test_reconcile_unit_level_partial_progress_on_mid_sequence_failure() -> None assert progress.purged_repos == [] +# --- streaming tarball ingestion (#106) ------------------------------------- + + +def _two_top_level_dirs_tarball() -> bytes: + """A real gzip tarball that ``indexer.ingest`` must REJECT. + + GitHub tarballs carry exactly one top-level ``org-repo-/`` directory; + two is the cheapest malformed shape that is still a structurally valid + archive, so the failure comes from ingestion rather than from gzip. + """ + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + for name in ("one/a.py", "two/b.py"): + info = tarfile.TarInfo(name) + info.size = len(b"x = 1\n") + tf.addfile(info, io.BytesIO(b"x = 1\n")) + return buf.getvalue() + + +class _ConnectCountingEngine(_FakeEngine): + """``_FakeEngine`` that counts ``connect()`` calls. + + ``run()`` takes exactly one connection outside the workers (the pre-fan-out + stamp read; the reconciliation checkpoint is gated off by any failure), so on + a one-repo failing run the count is 1 plus whatever the branch itself took. + """ + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.connects = 0 + + def connect(self) -> _FakeConn: + self.connects += 1 + return super().connect() + + +@pytest.mark.unit +def test_no_file_is_written_beside_the_tarball(monkeypatch: pytest.MonkeyPatch) -> None: + """AC2: the worker's temp dir holds the compressed tarball and nothing else. + + Asserted from inside ``index_fn`` -- i.e. while the branch is mid-flight and + an extracted tree would still exist -- not after the ``TemporaryDirectory`` + has already torn itself down. The spy is on ``indexer.job.download_tarball``: + ``job.py`` binds that name into its own namespace at import, so patching + ``indexer.fetch.download_tarball`` would miss. + """ + import indexer.job as job + + real_download_tarball = job.download_tarball + dests: list[Path] = [] + + def _download_tarball(*args: Any, **kwargs: Any) -> Any: + out = real_download_tarball(*args, **kwargs) + dests.append(out.parent) + return out + + monkeypatch.setattr(job, "download_tarball", _download_tarball) + + snapshots: list[list[str]] = [] + + def _index(conn: Any, *, items: Any, **_: Any) -> IndexCounts: + consumed = list(items) + snapshots.append(sorted(p.name for p in dests[-1].iterdir())) + return IndexCounts(files=len(consumed), symbols=0, swept=0, edges=0) + + assert _run(_config(repos=["acme/widgets"]), _index) == 0 + assert snapshots == [["source.tar.gz"]] + + +@pytest.mark.unit +def test_malformed_archive_fails_the_branch_from_inside_the_transaction( + caplog: pytest.LogCaptureFixture, +) -> None: + """D8, semantic-OFF path: archive validation moved inside the open transaction. + + ``extract_tarball`` raised before ``engine.connect()``. The lazy generator is + now consumed inside ``index_repo``'s ``conn.begin()``, so a malformed archive + costs one pooled connection and a rolled-back transaction. Accepted because + the branch-level OUTCOME is unchanged -- ``status="failed"``, exit code 1, no + ``phase timing`` line -- and the blast radius is one worker's own connection + (``pool_size == workers``, ``max_overflow=0``). + """ + github = _GitHub(tarball_bytes=_two_top_level_dirs_tarball()) + engine = _ConnectCountingEngine() + idx = _RecordingIndex() + + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + _config(repos=["acme/widgets"]), + idx, + cfg=Settings(semantic_enabled=False), + github=github, + engine=engine, + ) + + assert code == 1 + assert idx.counts == [] # the transaction never completed + assert engine.connects == 2 # the pre-fan-out stamp read, plus this branch's + assert any("failed to index acme/widgets@main" in r.getMessage() for r in caplog.records) + assert _timing_messages(caplog) == [] # a failed branch emits no timing line + + +@pytest.mark.unit +def test_malformed_archive_on_the_semantic_path_raises_before_any_connection() -> None: + """D8's twin: the PRODUCTION path still fails with zero connections held. + + On the semantic path ``job.py`` materializes the file list eagerly, so the + archive is fully streamed and fully validated before #104's advisory + ``shas_fn`` connection is opened and before ``index_repo``'s. The stamp below + matches ``INDEX_SEMANTICS_VERSION`` with a stale SHA, which is exactly the + shape that OPENS the delta gate -- so ``shas_fn`` would be called if the raise + were even one step later. + """ + github = _GitHub(tarball_bytes=_two_top_level_dirs_tarball()) + engine = _ConnectCountingEngine( + stamps={("acme/widgets", "main"): ("stale_sha", INDEX_SEMANTICS_VERSION)} + ) + idx = _RecordingIndex() + shas_calls: list[tuple[str, str]] = [] + + def _shas_fn(conn: Any, *, name: str, branch: str) -> tuple[set[Any], set[Any]]: + shas_calls.append((name, branch)) + return set(), set() + + code = _run( + _config(repos=["acme/widgets"]), + idx, + cfg=Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100), + embed_fn=lambda texts: [[0.0] for _ in texts], + github=github, + engine=engine, + shas_fn=_shas_fn, + ) + + assert code == 1 + assert idx.counts == [] + assert shas_calls == [] + assert engine.connects == 1 # the pre-fan-out stamp read alone + + # --- per-phase timing instrumentation (#103) -------------------------------- # Every INDEXED branch emits one `phase timing` record accounting for its whole # wall clock. The numbers are asserted EXACTLY, never with sleeps: the tests @@ -2634,7 +2782,7 @@ def test_reconcile_unit_level_partial_progress_on_mid_sequence_failure() -> None _TIMING_RE = re.compile( r"^phase timing (?P[^ @]+)@(?P\S+): " r"total=(\d+\.\d\d)s resolve=(\d+\.\d\d)s download=(\d+\.\d\d)s " - r"extract=(\d+\.\d\d)s parse=(\d+\.\d\d)s embed=(\d+\.\d\d)s " + r"parse=(\d+\.\d\d)s embed=(\d+\.\d\d)s " r"db=(\d+\.\d\d)s sweep=(\d+\.\d\d)s other=(\d+\.\d\d)s$" ) @@ -2679,7 +2827,7 @@ def _only_phases(caplog: pytest.LogCaptureFixture) -> dict[str, float]: @pytest.mark.unit def test_indexed_branch_logs_every_phase_field(caplog: pytest.LogCaptureFixture) -> None: - """T1: one record per indexed branch, with all nine fields in a fixed order.""" + """T1: one record per indexed branch, with all eight fields in a fixed order.""" idx = _RecordingIndex() with caplog.at_level(logging.INFO, logger="indexer.job"): code = _run(_config(repos=["acme/widgets"]), idx) @@ -2695,7 +2843,6 @@ def test_indexed_branch_logs_every_phase_field(caplog: pytest.LogCaptureFixture) "total", "resolve", "download", - "extract", "parse", "embed", "db", @@ -3165,25 +3312,25 @@ def test_semantic_eager_walk_is_charged_to_parse( ) -> None: """T18: on the semantic (production) path the file list is materialized up front, OUTSIDE the db window and outside _timed_items. Leaving that walk unwrapped would - dump the whole rglob + stat + read + decode cost into `other` -- the exact number + dump the whole stream + decompress + decode cost into `other` -- the exact number that routes work between the epic's parse/IO/db issues.""" import indexer.job as job clock = _install_fake_clock(monkeypatch) - real_iter_source_files = job.iter_source_files + real_iter_tar_source_files = job.iter_tar_source_files - def _slow_walk(root: Any) -> Any: - # iter_source_files is a generator function: paying the cost at call + def _slow_walk(tar_path: Any) -> Any: + # iter_tar_source_files is a generator function: paying the cost at call # time (before the first `next()`) would let this fake pass even if # the real wrap only timed the call and not the iteration -- the # exact mis-scoped-wrap bug this test exists to catch. Pay per - # yielded file instead, matching where the real cost (rglob walk, - # stat, read, decode) is actually incurred. - for pf in real_iter_source_files(root): + # yielded file instead, matching where the real cost (tar stream, + # decompress, decode) is actually incurred. + for pf in real_iter_tar_source_files(tar_path): clock.advance(3.0) # 2 files in _DEFAULT_FILES -> 6.0s total yield pf - monkeypatch.setattr(job, "iter_source_files", _slow_walk) + monkeypatch.setattr(job, "iter_tar_source_files", _slow_walk) cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) with caplog.at_level(logging.INFO, logger="indexer.job"): @@ -3238,7 +3385,7 @@ def test_other_is_the_exact_unattributed_residual( `other` is *defined* as `total - sum(phases)`, so asserting `total == sum(phases) + other` is a tautology that holds for any values and can never detect a mis-measured phase. This test instead drives a distinct known - amount through every one of the seven phases and requires `other == 0.00s`, + amount through every one of the six phases and requires `other == 0.00s`, then re-runs with one deliberately UNINSTRUMENTED advance and requires `other` to equal exactly that amount. """ @@ -3247,8 +3394,7 @@ def test_other_is_the_exact_unattributed_residual( clock = _install_fake_clock(monkeypatch) real_resolve_branch_head = job.resolve_branch_head real_download_tarball = job.download_tarball - real_extract_tarball = job.extract_tarball - real_iter_source_files = job.iter_source_files + real_iter_tar_source_files = job.iter_tar_source_files real_assert_disk_headroom = job.assert_disk_headroom def _resolve_branch_head(*args: Any, **kwargs: Any) -> str: @@ -3259,15 +3405,11 @@ def _download_tarball(*args: Any, **kwargs: Any) -> Any: clock.advance(2.0) return real_download_tarball(*args, **kwargs) - def _extract_tarball(*args: Any, **kwargs: Any) -> Any: - clock.advance(4.0) - return real_extract_tarball(*args, **kwargs) - - def _iter_source_files(root: Any) -> Any: - # Pay per yielded file, not at call time -- iter_source_files is a + def _iter_tar_source_files(tar_path: Any) -> Any: + # Pay per yielded file, not at call time -- iter_tar_source_files is a # generator function, so a call-time advance would pass even for a # wrap that only times the call and not the iteration. - for pf in real_iter_source_files(root): + for pf in real_iter_tar_source_files(tar_path): clock.advance(4.0) # 2 files in _DEFAULT_FILES -> 8.0s total yield pf @@ -3286,8 +3428,7 @@ def _index(conn: Any, *, items: Any, **_: Any) -> IndexCounts: monkeypatch.setattr(job, "resolve_branch_head", _resolve_branch_head) monkeypatch.setattr(job, "download_tarball", _download_tarball) - monkeypatch.setattr(job, "extract_tarball", _extract_tarball) - monkeypatch.setattr(job, "iter_source_files", _iter_source_files) + monkeypatch.setattr(job, "iter_tar_source_files", _iter_tar_source_files) cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) # "main" is stamped current so it is SKIPPED, leaving exactly one indexed @@ -3305,10 +3446,9 @@ def _index(conn: Any, *, items: Any, **_: Any) -> IndexCounts: phases = _only_phases(caplog) assert phases == { - "total": 127.0, + "total": 123.0, "resolve": 1.0, "download": 2.0, - "extract": 4.0, "parse": 8.0, "embed": 16.0, "db": 32.0, @@ -3330,10 +3470,9 @@ def _assert_disk_headroom(*args: Any, **kwargs: Any) -> None: phases = _only_phases(caplog) assert phases == { - "total": 136.0, + "total": 132.0, "resolve": 1.0, "download": 2.0, - "extract": 4.0, "parse": 8.0, "embed": 16.0, "db": 32.0, From 732a7d7ddbf7ea55bb74461fdfbf9e29acbe2af0 Mon Sep 17 00:00:00 2001 From: Tanner <84605639+IceRhymers@users.noreply.github.com> Date: Sat, 25 Jul 2026 09:50:33 -0700 Subject: [PATCH 5/9] test: watch indexer/ingest.py in the semantics tripwire (#106) Refs #106 --- docs/runbooks/indexing-parallelism.md | 4 +- indexer/store.py | 9 ++-- tests/unit/test_semantics_version_tripwire.py | 42 +++++++++++++++++++ 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index 892e0c3..042153c 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/docs/runbooks/indexing-parallelism.md @@ -427,8 +427,8 @@ no in-repo record of it. ## 5. Changing extraction semantics If you change **what** gets extracted — `indexer/symbols.py`, -`indexer/parse.py`, `indexer/languages.py` — you **must** bump -`INDEX_SEMANTICS_VERSION` in `app/db/models.py`. +`indexer/parse.py`, `indexer/languages.py`, `indexer/ingest.py` — you **must** +bump `INDEX_SEMANTICS_VERSION` in `app/db/models.py`. **The same obligation now extends past the tripwire's watched files (#104).** `indexer/parse.py`'s chunker is already a watched path, so a change to diff --git a/indexer/store.py b/indexer/store.py index 17f8c63..c2676b2 100644 --- a/indexer/store.py +++ b/indexer/store.py @@ -239,9 +239,12 @@ class is unioned in by ONE batched ``UPDATE ... RETURNING`` after the Two columns are deliberately NOT re-derived for a skipped file: * ``lang`` and ``size`` are pure functions of ``(path, content)`` via - ``indexer/parse.py`` + ``indexer/languages.py``, both watched by - ``tests/unit/test_semantics_version_tripwire.py`` -- so a change to either - derivation is MEANT to force a version bump, which closes this gate. That + ``indexer/ingest.py`` (the production file source) and + ``indexer/languages.py``'s ``EXT_TO_LANG`` (which it calls directly); + ``indexer/parse.py`` derives the same values in its retained role as + ``ingest.py``'s executable oracle. All three are watched by + ``tests/unit/test_semantics_version_tripwire.py`` -- so a change to any + of them is MEANT to force a version bump, which closes this gate. That tripwire is a local-developer guard rather than a CI one, so treat this as a strong convention backed by review, not a machine-enforced invariant. * ``files.commit`` goes staler. No production read path exists (every diff --git a/tests/unit/test_semantics_version_tripwire.py b/tests/unit/test_semantics_version_tripwire.py index 213d524..57c72bd 100644 --- a/tests/unit/test_semantics_version_tripwire.py +++ b/tests/unit/test_semantics_version_tripwire.py @@ -17,6 +17,40 @@ (shallow clones, tarball checkouts, a worktree with no remote): a tripwire that fails spuriously on a developer laptop gets disabled, and a disabled tripwire guards nothing. + +``indexer/ingest.py`` (#106) joined the watch set once it became the +production file source -- streaming tarball ingestion replaced +``indexer/parse.py``'s ``iter_source_files`` as the caller of the extraction +filter chain, and ``ingest.py`` also owns the ``lang``/``size`` derivation +that ``indexer.store``'s delta gate relies on the tripwire to protect +(``indexer/store.py``, the "two columns are deliberately NOT re-derived" +note). Leaving it unwatched would point that part of the guard at dead code. +``indexer/parse.py`` stays watched too: it still owns the production chunker +(``iter_chunks``, called from ``indexer/job.py``) and the shared binary-sniff +helper ``ingest.py`` imports, on top of its retained role as ``ingest.py``'s +executable oracle -- only ``iter_source_files`` lost its production caller. + +**Expected false positive, not confined to a single future event.** Adding a +new path to ``SEMANTICS_PATHS`` makes it an offender in any diff whose base +predates the file's creation -- ``git diff --name-only`` lists ADDED files, +not just changed ones. ``indexer/ingest.py`` postdates ``master`` (it was +added by #106's first PR, which has not merged past +``integration/indexer-performance``), so **every local run of this test on +this branch or its descendants sees it as an offender today** -- not only +the future diff that folds ``integration/indexer-performance`` into +``master`` (#111). ``_base_ref()`` falls back to ``origin/master`` locally +(there is usually no ``GITHUB_BASE_REF``), so this is the common case, not +the rare one; CI is actually less likely to see it, since both workflows +checkout at the default depth-1 and this test typically skips there instead +of running. That is expected, not a regression, and not a reason to disable +the test: resolve it via this tripwire's own documented escape below ("if +this change genuinely cannot alter extraction output, say so in the PR AND +ADJUST SEMANTICS_PATHS") -- the "say so in the PR" half only. Do NOT bump +``INDEX_SEMANTICS_VERSION`` and do NOT remove ``indexer/ingest.py`` from +``SEMANTICS_PATHS`` to silence it; the parity test +(``tests/unit/test_ingest_parity.py``, which pins ``ingest.py`` against +``parse.py`` field-for-field on the ``(path, lang, size, content)`` set) is +what actually backs the "no output change" claim. """ from __future__ import annotations @@ -32,10 +66,18 @@ # Modules that decide WHAT ends up in the index. A change to any of them can # alter extraction output for an unchanged commit, which is exactly the # condition the stored version stamp exists to detect. +# +# ``indexer/parse.py``'s ``iter_source_files`` keeps no production caller +# after #106, but the module stays watched: it still owns the production +# chunker (``iter_chunks``) and the shared binary sniff, on top of being +# ``indexer/ingest.py``'s executable oracle. See the module docstring above +# for why ``indexer/ingest.py`` joined this set and for the expected +# same-branch false positive. SEMANTICS_PATHS = ( "indexer/symbols.py", "indexer/parse.py", "indexer/languages.py", + "indexer/ingest.py", ) # The constant lives in app/db/models.py (it is read by both indexer.store, which From 86f8e8b0a9383dd9a57d80e6666f015e6965ee4c Mon Sep 17 00:00:00 2001 From: Tanner <84605639+IceRhymers@users.noreply.github.com> Date: Sat, 25 Jul 2026 11:18:11 -0700 Subject: [PATCH 6/9] indexer: batch per-file DB statements in index_repo (#105) Refs #105 --- docs/perf/issue-105-measurements.md | 318 +++++++++++ docs/runbooks/indexing-parallelism.md | 51 +- indexer/AGENTS.md | 14 +- indexer/bulk.py | 60 +++ indexer/chunk_store.py | 107 +++- indexer/job.py | 50 +- indexer/store.py | 347 ++++++++---- tests/integration/AGENTS.md | 4 +- tests/integration/test_chunk_batching.py | 298 +++++++++++ tests/integration/test_reconcile.py | 18 +- tests/integration/test_store_batching.py | 490 +++++++++++++++++ tests/integration/test_store_chunk_writer.py | 27 +- tests/unit/AGENTS.md | 6 +- tests/unit/test_bulk.py | 86 +++ tests/unit/test_chunk_store.py | 87 +++- tests/unit/test_job.py | 15 +- tests/unit/test_store_batching.py | 522 +++++++++++++++++++ tests/unit/test_store_chunk_writer.py | 47 +- tests/unit/test_store_delta.py | 55 +- 19 files changed, 2395 insertions(+), 207 deletions(-) create mode 100644 docs/perf/issue-105-measurements.md create mode 100644 indexer/bulk.py create mode 100644 tests/integration/test_chunk_batching.py create mode 100644 tests/integration/test_store_batching.py create mode 100644 tests/unit/test_bulk.py create mode 100644 tests/unit/test_store_batching.py diff --git a/docs/perf/issue-105-measurements.md b/docs/perf/issue-105-measurements.md new file mode 100644 index 0000000..780285e --- /dev/null +++ b/docs/perf/issue-105-measurements.md @@ -0,0 +1,318 @@ +# Issue #105 — batched writes: measurements + +Recorded once, at the point step 7 of the execution plan committed it, against +local Postgres 16 (`codesearch-pg`) and a throwaway `git worktree` at the +branch point `732a7d7` (`origin/integration/indexer-performance`, pre-#105). +Both scripts are reproduced here in full so the numbers below are re-derivable, +not just asserted; they are throwaway measurement scripts, not part of the test +suite (which is where the gate-asserted numbers — test 17's statement count and +test 14/15/15b/18's parity assertions — actually live). + +## 1. Statement count per file (AC 1: 3–7 → ≲0.05) + +Gate-asserted by `tests/integration/test_store_batching.py`'s +`test_statement_count_per_file_meets_the_acceptance_criterion` (test 17), a real +`before_cursor_execute` listener over a 600-file first-time index: + +``` +statements/file at _BATCH_MAX_FILES=500, N=600: 0.0133 (issue AC: <= 0.05) +``` + +Matches §2.1's predicted arithmetic (`16 statements / 500 files ≈ 0.032` at a +realistic symbol/edge mix; `0.02` on the pure per-file-only shape this fixture +uses) — measured, not just predicted, and comfortably inside the ≲0.05 target. + +## 2. Cross-version corpus parity (AC 3) + +The check `tests/integration/test_store_batching.py`'s tests 14/15/15b/18 +structurally cannot do: a diff against the *actual* pre-#105 implementation, +not just batched-against-batched. A 3-branch, 33-distinct-row fixture (5 +content-deduped files, 2 divergent-content files × 3 branches, one +zero-symbols file, one file with a real edge, 18 branch-unique files) indexed +identically in a throwaway schema on both trees, then dumped as +`(repo_id, path, content_sha, lang, size, content, commit, sorted(branches))` +for `files`, `(path, name, kind, start_line, end_line)` for `symbols`, and the +9-field tuple for `reference_edges` (serial ids excluded): + +``` +$ git worktree add /tmp/dcs105-perf/base-732a7d7 732a7d7 +$ cd /tmp/dcs105-perf/base-732a7d7 && uv run python parity_check.py before > before.json +$ cd && uv run python parity_check.py after > after.json +$ diff before.json after.json && echo "IDENTICAL — empty diff" +IDENTICAL — empty diff +``` + +`files: 33, symbols: 32, edges: 3` on both sides, dict-equal. The parity +harness script (`parity_check.py`, reproduced below) feeds `index_repo` an +`items` list directly, so it is unaffected by #106's change of file source — +exactly the reasoning §3.4 requires for comparing against `732a7d7`, not +`17aeb4f`. + +
+parity_check.py + +```python +import json +import sys +from uuid import uuid4 + +from sqlalchemy import text + +from app.db.client import create_db_engine +from app.db.models import Base +from indexer.languages import ExtractedEdge, ExtractedSymbol, FileExtraction, ParsedFile +from indexer.store import index_repo + + +def _pf(path, content): + return ParsedFile(path=path, lang="python", size=len(content.encode()), content=content) + + +def _sym(prefix, n): + return ExtractedSymbol(f"{prefix}{n}", "function", 1, 2) + + +def _fixture_items(branch): + items = [] + for i in range(5): + content = f"def shared{i}():\n return {i}\n" + items.append( + (_pf(f"shared{i}.py", content), FileExtraction(symbols=[_sym("shared", i)], edges=[])) + ) + for i in range(2): + magic = (sum(ord(c) for c in branch) * 31 + i) % 997 + content = f"def divergent{i}():\n return {magic}\n" + items.append( + ( + _pf(f"divergent{i}.py", content), + FileExtraction(symbols=[_sym("divergent", i)], edges=[]), + ) + ) + items.append((_pf("nosymbols.py", "# just a comment\n"), FileExtraction(symbols=[], edges=[]))) + caller = _sym("caller", 0) + items.append( + ( + _pf(f"{branch}_withedges.py", "def caller0():\n callee()\n"), + FileExtraction( + symbols=[caller], + edges=[ExtractedEdge(kind="call", target="callee", line=2, enclosing=caller)], + ), + ) + ) + for i in range(6): + content = f"def {branch}_only{i}():\n return {i}\n" + items.append( + ( + _pf(f"{branch}_only{i}.py", content), + FileExtraction(symbols=[_sym(f"{branch}_only", i)], edges=[]), + ) + ) + return items + + +def main(): + label = sys.argv[1] + schema = f"parity_{label}_{uuid4().hex[:8]}" + engine = create_db_engine() + conn = engine.connect() + try: + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + conn.execute(text(f"CREATE SCHEMA {schema}")) + conn.execute(text(f"SET search_path TO {schema}, public")) + conn.commit() + Base.metadata.create_all(bind=conn) + conn.commit() + + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a", + items=_fixture_items("a"), + ) + conn.commit() + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b", + items=_fixture_items("b"), + ) + conn.commit() + index_repo( + conn, + name="acme/widgets", + branch="c", + is_default=False, + head_sha="sha_c", + items=_fixture_items("c"), + ) + conn.commit() + + files = sorted( + tuple(r) + for r in conn.execute( + text( + "SELECT repo_id, path, content_sha, lang, size, content, commit, " + "array_to_string((SELECT array_agg(x ORDER BY x) FROM unnest(branches) x), ',') " + "FROM files ORDER BY path, content_sha" + ) + ).all() + ) + symbols = sorted( + tuple(r) + for r in conn.execute( + text( + "SELECT f.path, s.name, s.kind, s.start_line, s.end_line " + "FROM symbols s JOIN files f ON f.id = s.file_id" + ) + ).all() + ) + edges = sorted( + tuple(r) + for r in conn.execute( + text( + "SELECT f.path, e.edge_kind, e.target_name, e.line, e.enclosing_name, e.enclosing_kind, " + "e.enclosing_start_line, e.enclosing_end_line " + "FROM reference_edges e JOIN files f ON f.id = e.file_id" + ) + ).all() + ) + print(json.dumps({"files": files, "symbols": symbols, "edges": edges})) + finally: + conn.rollback() + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + conn.commit() + conn.close() + engine.dispose() + + +if __name__ == "__main__": + main() +``` + +
+ +## 3. `db=` before/after (AC 2: ≥5x on a full index) — NOT verified against Lakebase + +No dev Lakebase target was available in this environment, so per §3.4 this AC +is reported with both caveats and a labelled projection and is **not ticked**. +3000 first-time (full-path) files, `PhaseTimer` installed the same way +`indexer.job` installs it around `index_repo`, `db = wall − sweep` (no `parse` +term: `items` is a plain list already in hand, not lazily produced), three +runs per tree: + +| Tree | Run 1 | Run 2 | Run 3 | Mean | +|---|---|---|---|---| +| `732a7d7` (before, unbatched) | 4.1906s | 4.1697s | 4.2105s | **4.2036s** | +| this branch (after, batched) | 0.6550s | 0.6622s | 0.6607s | **0.6593s** | + +Measured local speedup: **4.2036 / 0.6593 ≈ 6.38×** — already past the issue's +5x threshold, even on loopback Postgres. Two caveats, both named up front +(§3.4), pushing in opposite directions: + +- *Understating:* local Postgres over TCP loopback has ~0.05–0.15 ms round + trips; Lakebase's are 1–2 ms. The protocol-chatter component of the win + (statement count × round-trip latency) is understated here by roughly an + order of magnitude relative to production. +- *Not understating:* the missing `symbols.file_id` index (§2.0 of the plan) + means the per-file/per-batch `DELETE FROM symbols WHERE file_id = ...` is a + sequential scan of the whole `symbols` table either way — CPU/IO-bound, not + round-trip-bound, so it shows up in this local number at full weight and is + not an artifact of the loopback environment. + +**Projection** (labelled as such, not measured): unbatched, this fixture's +plain per-file symbol-only shape costs 4 statements/file (file-upsert, +symbols-delete, symbols-insert, edges-delete); batched, it costs 0.0133 (§1, +test 17's 600-file gate). Δround-trips/file ≈ 4 − 0.0133 ≈ 3.99. At a Lakebase +round trip of 1–2 ms, that projects an ADDITIONAL 4.0–8.0 ms/file of pure protocol +chatter saved over the local number, on top of whatever the CPU-bound +`symbols` scan component already contributes locally — i.e. the production win +is expected to be **larger** than 6.38×, not smaller, but this is a projection, +not a Lakebase measurement, and AC 2 is carried as an open item on epic #110 +pending a real Lakebase run. + +
+timing_check.py + +```python +import sys +import time +from uuid import uuid4 + +from sqlalchemy import text + +from app.db.client import create_db_engine +from app.db.models import Base +from indexer.languages import ExtractedSymbol, FileExtraction, ParsedFile +from indexer.store import index_repo +from indexer.timing import PhaseTimer, install_timer, reset_timer + + +def _pf(path, content): + return ParsedFile(path=path, lang="python", size=len(content.encode()), content=content) + + +def _fixture(n): + items = [] + for i in range(n): + content = f"def f{i}():\n x = {i}\n return x\n" + items.append( + ( + _pf(f"pkg/mod{i}.py", content), + FileExtraction(symbols=[ExtractedSymbol(f"f{i}", "function", 1, 3)], edges=[]), + ) + ) + return items + + +def main(): + label = sys.argv[1] + n = int(sys.argv[2]) + schema = f"timing_{label}_{uuid4().hex[:8]}" + engine = create_db_engine() + conn = engine.connect() + try: + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + conn.execute(text(f"CREATE SCHEMA {schema}")) + conn.execute(text(f"SET search_path TO {schema}, public")) + conn.commit() + Base.metadata.create_all(bind=conn) + conn.commit() + + items = _fixture(n) + + timer = PhaseTimer() + token = install_timer(timer) + try: + wall_start = time.perf_counter() + index_repo( + conn, + name="acme/bigrepo", + branch="main", + is_default=True, + head_sha="sha1", + items=items, + ) + wall = time.perf_counter() - wall_start + finally: + reset_timer(token) + + sweep = timer.total("sweep") + db = wall - sweep + print(f"n={n} wall={wall:.4f}s sweep={sweep:.4f}s db(wall-sweep)={db:.4f}s") + finally: + conn.rollback() + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + conn.commit() + conn.close() + engine.dispose() + + +if __name__ == "__main__": + main() +``` + +
diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index 042153c..c7e99c7 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/docs/runbooks/indexing-parallelism.md @@ -157,7 +157,7 @@ promise that the field set never changes across releases.) | `download` | archive I/O bound | — (the decompression it used to be paired with is now fused into `parse`, #106) | | `parse` | GIL-bound tree-sitter extraction — **plus, since #106, the archive's gzip decompression, tar-stream read and UTF-8 decode**, which used to be the separate `extract=` field | #108 (process-pool extraction) addresses the tree-sitter half only | | `embed` | serial AI Gateway round trips | #107 (concurrent embedding) | -| `db` | per-file round trips | #105 (batched writes) | +| `db` | round trips for the batch's changed/new files, plus one Postgres-side sequential scan per file for its symbol delete (see §2.3) | #105 (batched writes) — landed; §2.3 | | any of the above, on **unchanged** content | redundant work | #104 (file-level delta indexing) | `#104` narrows the **db** and **embed** costs to a branch's actual delta, not @@ -242,6 +242,55 @@ check its `repo_branches.index_semantics_version` against the current `INDEX_SEMANTICS_VERSION` and whether a sibling branch is stale (§4's provenance gate). +### 2.3 Batched writes (#105) + +`index_repo`'s changed/new file loop no longer issues 5–7 statements per file +(a file upsert, a symbol delete-then-conditional-insert, an edge +delete-then-conditional-insert, and an optional chunk write). It accumulates +files into a batch and flushes once the batch reaches whichever bound trips +first — `_BATCH_MAX_FILES` (500 files) or `_BATCH_MAX_CONTENT_BYTES` (8 MiB of +`pf.size`) — plus once more after the loop for whatever remains. Each flush +issues ONE multi-row `files` upsert, ONE `DELETE ... WHERE file_id = ANY(...)` +each for `symbols`/`reference_edges` over the whole batch, a param-budgeted +bulk insert for each, and ONE `chunk_writer` call for the batch. The +membership-only class (already a single `UPDATE ... RETURNING` per branch) +gained nothing new to batch — its own chunk write collapsed from one call per +file to one call for the whole class as part of the same change. This is +invisible in the corpus: a batched write and a per-file write produce a +byte-identical `files`/`symbols`/`reference_edges` corpus (verified across +`_BATCH_MAX_FILES` ∈ {1, 2, 7, 500} and against the pre-#105 implementation +directly) — there is no new log line, and no existing line's format changed. + +**`db=` on the `phase timing` line is the observable.** Round trips per +changed/new file drop from 3–7 to a fraction of one; on a branch that is +mostly re-writing content (a first index, or the first full re-index after an +`INDEX_SEMANTICS_VERSION` bump), that collapses most of `db=`'s protocol-chatter +component. It does **not** collapse `db=` to zero: one component of `db=` is +CPU/IO-bound, not round-trip-bound, and batching does not touch it — `symbols` +carries no index on `file_id` (`app/db/models.py`; `reference_edges` has +`ix_reference_edges_file_id`, `symbols` has nothing), so `DELETE FROM symbols +WHERE file_id = ANY(...)` is still a sequential scan of the whole `symbols` +table per flush, just one scan per up-to-500 files instead of one per file. +Adding that index is tracked separately (a migration, deliberately kept out of +this change so it stays a zero-schema-change, `git revert`-safe deploy) and +would make this component cheaper too, independent of batch size. + +**The win concentrates on a first index or a post-bump full reindex.** Once a +branch is past its first run, file-level delta indexing (#104) already skips +`db` work entirely for unchanged content — batching only ever helps the +changed/new fraction that #104 leaves behind, so a branch with a small, +steady delta was already cheap and will not visibly move on this line. + +**The batch bounds are source constants, not a config knob.** `_BATCH_MAX_FILES` +and `_BATCH_MAX_CONTENT_BYTES` live in `indexer/store.py`, per this repo's +"guardrail constants with config-level fixes, not override flags" convention +(`indexer/AGENTS.md`). Raising `_BATCH_MAX_FILES` is a deploy, not a runtime +toggle, and is bounded by libpq's Bind-message parameter ceiling +(`_BATCH_MAX_FILES * _FILE_UPSERT_COLUMNS < 65535`, asserted at import time). +`_BATCH_MAX_FILES = 1` degrades the write path back to per-file statement +counts — useful for bisecting a suspected batching regression locally, but a +source edit and redeploy, never a 2am incident-response lever. + --- ## 3. The three limits, and why raising concurrency is a bad trade diff --git a/indexer/AGENTS.md b/indexer/AGENTS.md index f052873..268b661 100644 --- a/indexer/AGENTS.md +++ b/indexer/AGENTS.md @@ -11,7 +11,8 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` |------|-------------| | `__init__.py` | Empty package marker. | | `branches.py` | Resolves a connection's branch globs against a repo's real branch list via `fnmatchcase`, returning a typed `BranchResolution` (`branches`, `complete`, `dropped`, `cap`). Default branch always included and always first; empty globs means default-branch-only with no GitHub branches API call, always `complete=True`. `SOFT_BRANCH_CAP = 20` truncates (loud warning naming the cap and dropped branches, not a failure; sets `complete=False`, blocking reconciliation for that repo — #61); no override flag — fix the config. | -| `chunk_store.py` | `write_chunks`: delete-and-reinsert one file's rows in the `chunks` table (no natural key). Takes a live `Connection`; never opens an engine, never calls the embedder — vectors arrive precomputed. No `repo_id` parameter: `chunks` is scoped by `file_id` only. | +| `bulk.py` | `insert_rows` (#105): param-budgeted multi-row `INSERT` helper shared by `store.py` and `chunk_store.py` — explicit `.values([...])` rather than executemany, so the round-trip count is exactly one statement per slice on every driver. `PARAM_BUDGET` (30,000) and `CHUNK_PARAM_BUDGET` (6,000, payload- not param-bounded — each chunk row carries a 1024-float vector). No-op on an empty row list. | +| `chunk_store.py` | `write_chunks_batch` (#105): delete-and-reinsert MANY files' rows in the `chunks` table in one `DELETE ... WHERE file_id = ANY(:ids)` (over every id in the call, including zero-chunk files) plus one param-budgeted bulk insert via `bulk.insert_rows`. `write_chunks` is now a one-element wrapper over it (signature/return unchanged). Takes a live `Connection`; never opens an engine, never calls the embedder — vectors arrive precomputed. No `repo_id` parameter: `chunks` is scoped by `file_id` only. | | `fetch.py` | All GitHub HTTP: paginated org/user repo enumeration (`RepoMeta`), branch listing, `resolve_ref`/`resolve_branch_head` (branch name -> immutable SHA), streamed tarball download capped at `MAX_TARBALL_BYTES` (500 MB), `assert_disk_headroom` (0.5 GB per worker — the compressed tarball is the only artifact written to disk). No extraction: `MAX_EXTRACTED_BYTES` moved to `ingest.py` with #106 and is now a work cap, not a disk cap. `RateLimitError` is deliberately narrow: 429 always; 403 only with `Retry-After` or `X-RateLimit-Remaining: 0` — other 403s are permission failures. | | `ingest.py` | `iter_tar_source_files` (#106): the job's only file source. ONE forward pass over the downloaded tarball (`tarfile.open(mode="r:*")` driven by `tf.next()`; no `getmembers()`, no `extractall()`, nothing written to disk), yielding the same `ParsedFile` stream `parse.iter_source_files` would have yielded from an extracted tree — same filter chain, reusing `parse._looks_binary` by import so the binary-sniff rule cannot fork. `size` is `member.size` (the archive-declared length, which exceeds `len(content)` after NUL stripping); `lang` is the RELATIVE path's lowered suffix. Yields in ARCHIVE order (deterministic per commit; nothing downstream depends on the sequence — see the module docstring's D5 note). `tf.members.clear()` is the FIRST statement of the loop body, since the `TarFile` now outlives `engine.connect()`. Replaces `filter="data"` with explicit, probe-verified rules: raw-`..`-component and absolute member names raise, absolute/escaping link targets raise for hardlinks and symlinks alike, special files and benign hardlinks skip with a WARNING, duplicate names are first-wins, and a zero-member archive raises rather than silently stamping the branch as indexed. `MAX_EXTRACTED_BYTES` (2 GB) is enforced incrementally over pre-filter member sizes. | | `hashing.py` | `content_sha`: canonical SHA-256 hex of content (`None` -> empty string). Single source of truth for `files.content_sha`; must stay byte-identical to the `0003` migration's SQL backfill forever or cross-branch dedup silently breaks (`tests/integration/test_content_sha_parity.py` is the gate). | @@ -20,7 +21,7 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` | `parse.py` | `iter_source_files`: **no production caller since #106** — retained deliberately as the extraction ORACLE the streaming path (`ingest.py`) is pinned against by `tests/unit/test_ingest_parity.py`, and undeletable anyway without editing a `SEMANTICS_PATHS` module and forcing an `INDEX_SEMANTICS_VERSION` bump. Walk an extracted tree, yield every text file (unknown extensions kept with `lang=None`, since grep runs over all files); skips `.git/`, symlinks, files > `MAX_FILE_BYTES` (stat before read), NUL-sniffed binaries, and UTF-8 decode failures; strips surviving NULs (Postgres `text` rejects them). `iter_chunks`: deterministic line-aligned chunking, no overlap, no mid-line splits, 1-based inclusive line ranges. | | `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (0.5 GB peak per worker since #106), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. Deliberately import-light: pydantic + PyYAML + stdlib only. | | `resolve.py` | `resolve_repos`: enumerate org/user selectors, apply `ExcludeRules` to enumerated repos only (explicit `repos` entries always win, unfiltered), dedup case-insensitively keeping first-seen spelling, union each connection's `branches:` globs into per-repo `RepoEntry.branch_globs`, then fail fast: `EmptyConfigError` on zero repos (indexing nothing must not exit 0), `RepoCeilingError` above `MAX_REPOS` (500, overridable via `--max_repos`). Also matches `RepoConfig.semantic_max_chunks_per_repo` case-insensitively onto each resolved repo's `RepoEntry.semantic_max_chunks` (`None` if unmatched); an override key matching no resolved repo logs a WARNING (typo guard). Enumerator params are a test seam only, not provider dispatch. | -| `store.py` | `index_repo`: the single atomic unit of work for one (repo, branch) inside `with conn.begin():` — repos upsert, `repo_branches` CAS-baseline read, per-file array-union upsert on `(repo_id, path, content_sha)` with branch-membership union, delete-and-reinsert `symbols` AND `reference_edges` (both keyed only by `file_id`, no natural key), optional `chunk_writer` call, membership sweep (strip this branch from unseen rows, delete rows with empty `branches`; skipped with a WARNING on an empty seen-set), then CAS stamp (raises `StaleIndexError` on mismatch, rolling everything back). `items` is `Iterable[tuple[ParsedFile, FileExtraction]]`; the `reference_edges` delete runs unconditionally, even when a file's `FileExtraction.edges` is empty, so stale rows never survive a re-index. The sweep call site is wall-clocked into `indexer.timing`'s ambient `PhaseTimer` (`record("sweep", …)`) — never into `IndexCounts` (frozen, compared by value) and never through a new `index_repo` parameter (that signature is an injected seam); `store.py` emits **no** log record for it. Pure DML, no TEMP tables (job role has no TEMP privilege on Lakebase). `reconcile_retired_branches` and `reconcile_removed_repos`: the desired-state storage primitives `job.py`'s reconciliation checkpoint calls — each its own `conn.begin()`, repo-scoped membership subtraction (never delete-by-path) and an exact-match repo purge respectively; neither decides WHAT is retired/desired, only applies a caller-supplied set (see `job.py`'s row). | +| `store.py` | `index_repo`: the single atomic unit of work for one (repo, branch) inside `with conn.begin():` — repos upsert, `repo_branches` CAS-baseline read, per-file classification (unchanged / membership-only / changed-new), membership sweep (strip this branch from unseen rows, delete rows with empty `branches`; skipped with a WARNING on an empty seen-set), then CAS stamp (raises `StaleIndexError` on mismatch, rolling everything back). `items` is `Iterable[tuple[ParsedFile, FileExtraction]]`. **Changed/new files are batched (#105)**, not written one at a time: `_flush_file_batch` accumulates up to `_BATCH_MAX_FILES` files or `_BATCH_MAX_CONTENT_BYTES` bytes of `pf.size` (whichever trips first, post-append — a single oversized file is never split), then issues ONE multi-row `files` upsert (`SET` uses `excluded.*` for every per-row column, NEVER a Python literal from one file — the trap that would collapse a batch's conflicting rows onto the last file's values), maps ids back from `RETURNING` by `(path, content_sha)` — never by row order, and RAISES on a missing id rather than warn-and-skip (a wrong id here would attach one file's symbols to another file's row) — then ONE `DELETE ... WHERE file_id = ANY(:ids)` + bulk insert each for `symbols`/`reference_edges` (`reference_edges` unconditional, even when a batch's `FileExtraction.edges` is all-empty, so stale rows never survive a re-index), then ONE `chunk_writer` call for the batch. An intra-batch duplicate `(path, content_sha)` is deduped first (last occurrence wins, one WARNING) — Postgres raises `ON CONFLICT DO UPDATE command cannot affect row a second time` and poisons the transaction otherwise. The sweep call site is wall-clocked into `indexer.timing`'s ambient `PhaseTimer` (`record("sweep", …)`) — never into `IndexCounts` (frozen, compared by value) and never through a new `index_repo` parameter (that signature is an injected seam); `store.py` emits **no** log record for it. Pure DML, no TEMP tables (job role has no TEMP privilege on Lakebase). `reconcile_retired_branches` and `reconcile_removed_repos`: the desired-state storage primitives `job.py`'s reconciliation checkpoint calls — each its own `conn.begin()`, repo-scoped membership subtraction (never delete-by-path) and an exact-match repo purge respectively; neither decides WHAT is retired/desired, only applies a caller-supplied set (see `job.py`'s row). | | `symbols.py` | `extract_file` (#84): one tree-sitter parse, one full-tree walk, emitting both `symbols` and `reference_edges` candidates (`FileExtraction`) in the same pass — nested definitions captured, named nodes only for symbols, kinds from `SYMBOL_KINDS`/`EDGE_NODE_KINDS`, 1-based lines. Stack entries carry the innermost named enclosing definition so edges attribute to it in O(1) with no second walk. `extract_symbols` is a thin wrapper (`extract_file(pf).symbols`) kept for the existing call sites. Python-only edge helpers (`_python_call_target`, `_python_import_edges`) are the seam #85 generalizes to other languages. Parser cache is per-thread (`threading.local`) as insurance against a future GIL-releasing `parse()`. | | `timing.py` | Per-phase wall-clock accounting for one branch (#103). `PhaseTimer` accumulates `phase -> seconds`; `install_timer`/`reset_timer`/`current_timer` carry one ambiently in a `ContextVar` (the same idiom as `job.py`'s `_repo_ctx`, for the same cross-module attribution problem); `record(phase, seconds)` is a **no-op when no timer is installed and never raises**, so `index_repo` stays callable directly. `_CLOCK` (default `time.monotonic`) is the single clock source every asserted duration reads — via `now()`, or via `PhaseTimer.clock` captured at construction — and is the seam tests patch; never read `time.monotonic()` directly for a number that appears on those lines. Stdlib only. | @@ -33,15 +34,16 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - **One logical corpus writer, at the run level**: `resources/job.yml` pins `max_concurrent_runs: 1` (queueing retained), so at most one run of this job is ever active — the invariant global desired-state reconciliation depends on. This is separate from (and does not replace) the per-branch sequencing above: it bounds concurrent job RUNS, not per-run concurrency, and it does NOT cover a writer outside this job (a second job, a manual `bundle run`, or future per-repo/per-branch task sharding within one run) — any of those needs a shared database fencing/lease protocol first. See `docs/runbooks/indexing-parallelism.md` §1.1 for the full invariant and its coverage boundary. - **Reconciliation is a clean-run-only, corpus-wide, main-thread-only checkpoint**: it reuses the pre-fan-out stamp snapshot (never re-reads `repo_branches` after fan-out) — sound only because of the single-writer-run invariant above; do not relax `max_concurrent_runs` or add in-run task sharding without re-deriving that proof (see the loud comment at the checkpoint in `job.py`). Reconciliation must never be invoked from `_index_one`/`_index_one_inner`/`_index_one_branch`, and its two seams (`reconcile_retired_fn`, `reconcile_removed_fn`) must always default to the real `indexer.store` primitives, mirroring `index_fn`'s injection pattern. - **`content_sha` parity is forever**: `hashing.content_sha` must stay byte-identical to the `0003` migration's SQL expression, or content dedup silently mints duplicate rows. -- **No network inside the transaction**: embeddings are precomputed in `job._precompute_chunk_writer` before `engine.connect()`; `chunk_store.write_chunks` is pure DML. The semantic layer is additive — every semantic failure (embedder down, chunk ceiling, count mismatch) degrades to indexing the core corpus without chunks, never to losing the branch. The chunk ceiling itself is `entry.semantic_max_chunks or cfg.semantic_max_chunks_per_repo`, resolved once per repo in `_index_one_inner` and threaded through `_index_one_branch` as `max_chunks_per_repo` — never read from `cfg` directly inside the branch loop, so every branch of one repo enforces the same resolved cap. +- **No network inside the transaction**: embeddings are precomputed in `job._precompute_chunk_writer` before `engine.connect()`; `chunk_store.write_chunks`/`write_chunks_batch` are pure DML. The semantic layer is additive — every semantic failure (embedder down, chunk ceiling, count mismatch) degrades to indexing the core corpus without chunks, never to losing the branch. The chunk ceiling itself is `entry.semantic_max_chunks or cfg.semantic_max_chunks_per_repo`, resolved once per repo in `_index_one_inner` and threaded through `_index_one_branch` as `max_chunks_per_repo` — never read from `cfg` directly inside the branch loop, so every branch of one repo enforces the same resolved cap. +- **`ChunkWriter` is called once per BATCH, not once per file (#105)**: `Callable[[Connection, int, Sequence[tuple[int, ParsedFile]]], None]` — `store.py`'s `_flush_file_batch` and `_union_membership` each call it once with the whole batch/membership-class's `(file_id, pf)` pairs. `job._precompute_chunk_writer`'s closure applies its `covered`-path guard per pair before one `write_chunks_batch` call for whatever survives. - **Redaction**: the GitHub token is read from Databricks secrets (base64), lives only in the injected `httpx.Client`'s `Authorization` header, and is never logged; request headers are never logged; the job never lowers root/SDK/httpx log levels (`tests/unit/test_job_redaction.py` enforces this). - **Test seams are injected callables**, not registries: `run()` accepts `workspace_client`, `http_client`, `engine`, `index_fn`, `embed_fn`, `config_loader`, `reconcile_retired_fn`, `reconcile_removed_fn`; `resolve_repos` accepts enumerators; `store.index_repo`/`reconcile_retired_branches`/`reconcile_removed_repos` and `chunk_store.write_chunks` take a live `Connection`. `resolve_repos` itself is deliberately NOT injectable in `run()` — tests drive it through `httpx.MockTransport`. - **The schema is provider-extensible; the code is GitHub-only.** No `type`-keyed dispatch anywhere — adding GitLab means a new union member plus a new fetch implementation, not a provider abstraction. - **Import discipline**: `repo_config.py` stays import-light (no httpx/SQLAlchemy/SDK at module level); `normalize_repo` is imported from `repo_config`, never from `job` (import cycle); `app.embed`'s databricks-sdk dependency is only imported when semantic is on. ### Testing Requirements -- `make test` (`pytest -m "unit or observability"`, no external deps): `tests/unit/test_branches.py`, `test_chunk_store.py`, `test_chunking.py`, `test_edges.py`, `test_fetch.py`, `test_ingest.py`, `test_ingest_parity.py`, `test_job.py`, `test_job_redaction.py`, `test_languages.py`, `test_parse.py`, `test_repo_config.py`, `test_resolve.py`, `test_symbols.py`, `test_store_chunk_writer.py`, `test_semantics_version_tripwire.py`, `test_timing.py`. -- `make test-integration` (needs Postgres): `tests/integration/test_store.py`, `test_store_chunk_writer.py`, `test_content_sha_parity.py` (the hard gate on `hashing.py` vs the migration backfill), `test_reconcile.py` (the `reconcile_retired_branches`/`reconcile_removed_repos` storage primitives directly), and `test_job_reconcile.py` (#59: the real `run()` end to end — real primitives, real engine, only GitHub HTTP faked). +- `make test` (`pytest -m "unit or observability"`, no external deps): `tests/unit/test_branches.py`, `test_bulk.py`, `test_chunk_store.py`, `test_chunking.py`, `test_edges.py`, `test_fetch.py`, `test_ingest.py`, `test_ingest_parity.py`, `test_job.py`, `test_job_redaction.py`, `test_languages.py`, `test_parse.py`, `test_repo_config.py`, `test_resolve.py`, `test_symbols.py`, `test_store_batching.py`, `test_store_chunk_writer.py`, `test_store_delta.py`, `test_semantics_version_tripwire.py`, `test_timing.py`. +- `make test-integration` (needs Postgres): `tests/integration/test_store.py`, `test_store_batching.py` (#105, batch-size parity/excluded-trap/rollback/statement-count — creates no `chunks` table, runs locally), `test_chunk_batching.py` (#105, batched `write_chunks_batch` on bare `vector`, runs locally), `test_store_chunk_writer.py`, `test_content_sha_parity.py` (the hard gate on `hashing.py` vs the migration backfill), `test_reconcile.py` (the `reconcile_retired_branches`/`reconcile_removed_repos` storage primitives directly), and `test_job_reconcile.py` (#59: the real `run()` end to end — real primitives, real engine, only GitHub HTTP faked). - Schema tests must stay fast: changes to `repo_config.py` must not add heavy imports. ### Common Patterns @@ -49,7 +51,7 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - `pg_insert(...).on_conflict_do_update(...).returning(...)` with a no-op `SET` on conflict — `DO NOTHING ... RETURNING` returns no row on conflict and would break the id bootstrap. - Delete-and-reinsert for child rows with no natural key (symbols, chunks). - `fnmatchcase`, never `fnmatch`: plain globs must behave identically on every platform. -- Guardrail constants with config-level fixes, not override flags: `SOFT_BRANCH_CAP` (20), `MAX_REPOS` (500, the one exception — `--max_repos`), `MAX_FILE_BYTES`, `MAX_TARBALL_BYTES` (`fetch.py`, a disk cap), `MAX_EXTRACTED_BYTES` (`ingest.py` since #106, a WORK cap — nothing is written to disk, so it bounds streamed content, not storage), `MAX_PURGE_SHRINK_FRACTION` (0.5, `job.py` — withholds a corpus-wide repo purge above this shrink fraction; fix is staged config removal, never a knob). +- Guardrail constants with config-level fixes, not override flags: `SOFT_BRANCH_CAP` (20), `MAX_REPOS` (500, the one exception — `--max_repos`), `MAX_FILE_BYTES`, `MAX_TARBALL_BYTES` (`fetch.py`, a disk cap), `MAX_EXTRACTED_BYTES` (`ingest.py` since #106, a WORK cap — nothing is written to disk, so it bounds streamed content, not storage), `MAX_PURGE_SHRINK_FRACTION` (0.5, `job.py` — withholds a corpus-wide repo purge above this shrink fraction; fix is staged config removal, never a knob), `_BATCH_MAX_FILES` (500) / `_BATCH_MAX_CONTENT_BYTES` (8 MiB, `store.py` since #105 — the changed/new write-batch bounds; `_BATCH_MAX_FILES * _FILE_UPSERT_COLUMNS < 65535` is asserted at import time against libpq's Bind param ceiling), `PARAM_BUDGET` (30,000) / `CHUNK_PARAM_BUDGET` (6,000, `bulk.py` since #105 — the generic multi-row-insert slicing bounds). - Errors name the fix: exceptions carry the config key or job parameter to change (`exclude.size_mb`, `index_concurrency`, `--max_repos`). - Broad `except Exception` at classification boundaries only (`run()`'s config/resolve handlers, `_index_one_branch`), each with a comment explaining why broad is load-bearing. - Lazy generators through the open transaction for bounded memory (semantic mode is the exception — it materializes the file list because embedding needs all chunk texts up front). diff --git a/indexer/bulk.py b/indexer/bulk.py new file mode 100644 index 0000000..05decbb --- /dev/null +++ b/indexer/bulk.py @@ -0,0 +1,60 @@ +"""Param-budgeted multi-row ``INSERT`` helper shared by ``store.py`` and ``chunk_store.py``. + +A separate module rather than a private helper in ``store.py``: ``chunk_store.py`` +needs it too and deliberately does not import ``store.py`` (its own docstring says +it mirrors that module's connection seam, not depends on it). Not a +``SEMANTICS_PATHS`` file -- it changes how rows are written, never what is +extracted. + +**Why an explicit multi-row ``.values([...])`` and not ``conn.execute(pg_insert(T), +[rows])``.** The executemany form is what the per-file code used before this +module existed, and SQLAlchemy would page it for us. But the round-trip +acceptance criterion this module exists to satisfy is stated in *statements*, and +executemany's round-trip count is a property of psycopg3's pipeline mode and +libpq's version -- an environment-dependent number. An explicit multi-row +``VALUES`` is unambiguously one statement on every driver and every target, and is +measurable identically in a unit fake, an integration test, and production. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from sqlalchemy import Connection +from sqlalchemy.dialects.postgresql import insert as pg_insert + +# libpq's Bind message carries the parameter count in an int16, so no single +# statement may bind more than 65535 params. Callers size their own batches +# against this budget (see indexer/store.py's _BATCH_MAX_FILES * +# _FILE_UPSERT_COLUMNS < 65535 invariant); this is the generic slicing bound +# for the tables that are not directly param-count-limited by a batch cap. +PARAM_BUDGET = 30_000 + +# Payload-bounded, not param-bounded: each chunk row carries a 1024-float +# embedding vector, so ~1000 rows/statement keeps the wire payload around 4 MB +# instead of 20 MB. +CHUNK_PARAM_BUDGET = 6_000 + + +def insert_rows( + conn: Connection, + table: Any, + rows: Sequence[dict[str, Any]], + *, + param_budget: int = PARAM_BUDGET, +) -> None: + """Issue one multi-row ``INSERT`` per ``param_budget``-sized slice of ``rows``. + + A no-op on an empty ``rows`` -- no statement is issued at all. Every row is + present exactly once, in input order; no single statement binds more than + ``param_budget`` params (``rows_per_statement = param_budget // + len(rows[0])``, so this also bounds each statement's row count). + """ + if not rows: + return + columns = len(rows[0]) + rows_per_statement = max(1, param_budget // columns) + for start in range(0, len(rows), rows_per_statement): + chunk = rows[start : start + rows_per_statement] + conn.execute(pg_insert(table).values(list(chunk))) diff --git a/indexer/chunk_store.py b/indexer/chunk_store.py index 3d077c2..7ae72ba 100644 --- a/indexer/chunk_store.py +++ b/indexer/chunk_store.py @@ -1,11 +1,11 @@ -"""Write PRECOMPUTED chunk+embedding rows for one file into ``chunks``. +"""Write PRECOMPUTED chunk+embedding rows for many files into ``chunks``. Mirrors ``indexer.store``'s connection seam: the caller supplies a live ``sqlalchemy.Connection`` and owns the transaction (``conn.begin()``); this module never opens its own engine. Like ``index_repo``'s symbol handling, -chunks carry no natural key, so a re-index deletes ``file_id``'s existing rows +chunks carry no natural key, so a re-index deletes each file's existing rows and reinserts the current set -- idempotent, and safe to call repeatedly -within the same per-file loop. +within the same per-batch flush. This module never calls the embedder: ``chunks`` arrives with vectors already computed by :mod:`app.embed`, so writing them is pure DML with no network @@ -21,43 +21,94 @@ from __future__ import annotations +import logging from collections.abc import Sequence -from sqlalchemy import Connection, delete -from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy import ARRAY, BigInteger, Connection, any_, bindparam, delete from app.db.semantic import chunks as chunks_table +from indexer.bulk import CHUNK_PARAM_BUDGET, insert_rows +logger = logging.getLogger("indexer.chunk_store") -def write_chunks( +# One (chunk_index, content, start_line, end_line, embedding) tuple per chunk, +# in the shape the line-aligned chunker produces. +ChunkRow = tuple[int, str, int, int, list[float]] + + +def write_chunks_batch( conn: Connection, *, - file_id: int, - chunks: Sequence[tuple[int, str, int, int, list[float]]], + rows: Sequence[tuple[int, Sequence[ChunkRow]]], ) -> int: - """Delete-and-reinsert ``file_id``'s chunk rows; return the row count written. + """Delete-and-reinsert chunk rows for many files in one delete + one bulk insert. - ``chunks`` is a sequence of ``(chunk_index, content, start_line, end_line, - embedding)`` tuples with embeddings already computed and 1-based inclusive - line ranges from the line-aligned chunker. Runs inside the caller's open - transaction, alongside the rest of that file's ``index_repo`` work. + ``rows`` is a sequence of ``(file_id, chunks)`` pairs. A zero-statement + no-op on an empty ``rows`` -- no ``conn`` call at all. Otherwise: one + ``DELETE ... WHERE file_id = ANY(:ids)`` over EVERY file id in the call -- + including files with zero chunk rows, since delete-on-zero-chunks is + load-bearing (a file that shrank to zero chunks must still lose its stale + rows) -- then one param-budgeted bulk insert via :func:`indexer.bulk.insert_rows`. + + **Dedup guard, keyed on ``file_id``, keeping the LAST occurrence** (mirrors + ``indexer.store._flush_file_batch``'s ``(path, content_sha)`` guard). + Mandatory here too: two rows for the same ``file_id`` would insert two + conflicting ``(file_id, chunk_index)`` tuples, violating + ``uq_chunks_file_id_chunk_index`` and poisoning the transaction. Callers + are expected not to duplicate a ``file_id`` in one call, but this is a + shared primitive reachable from more than one caller (``indexer/job.py``'s + changed/new-path closure, and ``indexer/store.py``'s ``_union_membership``, + which has no dedup guard of its own on its injected ``items`` seam), so the + guard lives here rather than being duplicated at every call site. """ - conn.execute(delete(chunks_table).where(chunks_table.c.file_id == file_id)) - if not chunks: + if not rows: return 0 + deduped: dict[int, Sequence[ChunkRow]] = {} + for file_id, chunks in rows: + if file_id in deduped: + logger.warning( + "duplicate file_id %r within one write_chunks_batch call; " + "keeping the last occurrence", + file_id, + ) + deduped[file_id] = chunks + entries = list(deduped.items()) + + ids = [file_id for file_id, _chunks in entries] conn.execute( - pg_insert(chunks_table), - [ - { - "file_id": file_id, - "chunk_index": chunk_index, - "content": content, - "start_line": start_line, - "end_line": end_line, - "embedding": embedding, - } - for chunk_index, content, start_line, end_line, embedding in chunks - ], + delete(chunks_table).where( + chunks_table.c.file_id == any_(bindparam("ids", type_=ARRAY(BigInteger))) + ), + {"ids": ids}, ) - return len(chunks) + + insert_dicts = [ + { + "file_id": file_id, + "chunk_index": chunk_index, + "content": content, + "start_line": start_line, + "end_line": end_line, + "embedding": embedding, + } + for file_id, chunks in entries + for chunk_index, content, start_line, end_line, embedding in chunks + ] + insert_rows(conn, chunks_table, insert_dicts, param_budget=CHUNK_PARAM_BUDGET) + return len(insert_dicts) + + +def write_chunks( + conn: Connection, + *, + file_id: int, + chunks: Sequence[ChunkRow], +) -> int: + """Delete-and-reinsert ``file_id``'s chunk rows; return the row count written. + + A one-element wrapper over :func:`write_chunks_batch` -- one DML + implementation, not two that can drift. Signature and return value are + unchanged from before batching. + """ + return write_chunks_batch(conn, rows=[(file_id, chunks)]) diff --git a/indexer/job.py b/indexer/job.py index b4ccb3f..c5c5c6f 100644 --- a/indexer/job.py +++ b/indexer/job.py @@ -125,7 +125,7 @@ import sys import tempfile import time -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Callable, Iterable, Iterator, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed from contextvars import ContextVar from dataclasses import dataclass, field @@ -140,7 +140,7 @@ from app.db.models import INDEX_SEMANTICS_VERSION, Repo, RepoBranch from app.embed import EmbeddingCountMismatchError, EmbedFn, get_embedder from indexer.branches import resolve_branches -from indexer.chunk_store import write_chunks +from indexer.chunk_store import write_chunks_batch from indexer.fetch import ( REQUIRED_FREE_BYTES, assert_disk_headroom, @@ -869,11 +869,17 @@ def _precompute_chunk_writer( file in the branch. The closure below therefore closes over ``covered = set(per_file)`` -- every path THIS call embedded, including zero-chunk files -- and refuses to write chunks for any other path. This is defence-in-depth - for a path the single-writer-per-repo invariant says is unreachable: ` - `write_chunks`` deletes a file's chunk rows before inserting, so calling it - for a path this precompute never embedded would silently delete that file's - chunks rather than merely leave them stale. See ``indexer/store.py``'s - ``_union_membership`` for the authoritative-side analogue of this guard. + for a path the single-writer-per-repo invariant says is unreachable: + ``write_chunks_batch`` deletes a file's chunk rows before inserting, so + calling it for a path this precompute never embedded would silently delete + that file's chunks rather than merely leave them stale. See + ``indexer/store.py``'s ``_union_membership`` for the authoritative-side + analogue of this guard. + + The returned closure is called once per BATCH of files (a sequence of + ``(file_id, pf)`` pairs, per :data:`indexer.store.ChunkWriter`) -- the + ``covered`` guard is applied per pair, filtering the batch before one + :func:`indexer.chunk_store.write_chunks_batch` call for whatever survives. """ per_file: dict[str, list[Chunk]] = {pf.path: list(iter_chunks(pf)) for pf in files} covered = set(per_file) @@ -905,19 +911,23 @@ def _precompute_chunk_writer( ] i += len(chunks) - def chunk_writer(conn: Any, repo_id: int, file_id: int, pf: ParsedFile) -> None: - if pf.path not in covered: - # Unreachable while the single-writer invariant holds (see the - # docstring): index_repo only ever calls chunk_writer for a file this - # same precompute either embedded or classified membership-only (and - # index_repo's own _union_membership guards that case separately). - # Warn-and-skip rather than raise, matching this module's established - # additive-layer posture for the semantic path. - logger.warning( - "no precomputed chunks for %s; leaving its chunk rows untouched", pf.path - ) - return - write_chunks(conn, file_id=file_id, chunks=by_path.get(pf.path, [])) + def chunk_writer(conn: Any, repo_id: int, pairs: Sequence[tuple[int, ParsedFile]]) -> None: + rows: list[tuple[int, list[tuple[int, str, int, int, list[float]]]]] = [] + for file_id, pf in pairs: + if pf.path not in covered: + # Unreachable while the single-writer invariant holds (see the + # docstring): index_repo only ever calls chunk_writer for a file + # this same precompute either embedded or classified + # membership-only (and index_repo's own _union_membership guards + # that case separately). Warn-and-skip rather than raise, + # matching this module's established additive-layer posture for + # the semantic path. + logger.warning( + "no precomputed chunks for %s; leaving its chunk rows untouched", pf.path + ) + continue + rows.append((file_id, by_path.get(pf.path, []))) + write_chunks_batch(conn, rows=rows) return chunk_writer diff --git a/indexer/store.py b/indexer/store.py index c2676b2..2bfe0ce 100644 --- a/indexer/store.py +++ b/indexer/store.py @@ -20,19 +20,50 @@ from __future__ import annotations import logging -from collections.abc import Callable, Collection, Iterable +from collections.abc import Callable, Collection, Iterable, Sequence from dataclasses import dataclass - -from sqlalchemy import Connection, delete, func, text, update +from typing import cast + +from sqlalchemy import ( + ARRAY, + BigInteger, + Connection, + Table, + any_, + bindparam, + delete, + func, + text, + update, +) from sqlalchemy.dialects.postgresql import insert as pg_insert from app.db.models import INDEX_SEMANTICS_VERSION, File, ReferenceEdge, Repo, RepoBranch, Symbol +from indexer.bulk import insert_rows from indexer.hashing import content_sha from indexer.languages import FileExtraction, IndexCounts, ParsedFile from indexer.timing import now, record logger = logging.getLogger("indexer.store") +# Batch bounds for the changed/new write path (#105). Guardrail constants, not +# Settings fields (indexer/AGENTS.md: "guardrail constants with config-level +# fixes, not override flags") -- raising _BATCH_MAX_FILES is a deploy, never a +# runtime knob. +_BATCH_MAX_FILES = 500 +# Bounds two things at once: (1) peak retention of this buffer, additive to +# ingest.py's MAX_FILE_BYTES + seen-path set (#106) -- this buffer did not +# exist before #106 either, since the pre-#106 extracted tree held the whole +# corpus on disk; (2) the wire payload of the single `files` multi-row INSERT, +# kept in the same order as CHUNK_PARAM_BUDGET's ~4 MB target. +_BATCH_MAX_CONTENT_BYTES = 8 * 1024 * 1024 +# repo_id, path, lang, size, content, commit, content_sha, branches. +_FILE_UPSERT_COLUMNS = 8 +# INVARIANT: libpq's Bind message carries the parameter count in an int16, so +# no single statement may bind more than 65535 params. 500 * 8 = 4000, 16x +# headroom. Asserted in tests/unit/test_store_batching.py. +assert _BATCH_MAX_FILES * _FILE_UPSERT_COLUMNS < 65535 + class StaleIndexError(RuntimeError): """The ``repo_branches`` row changed between this transaction's first and last statement. @@ -64,10 +95,11 @@ class ReconcileCounts: files_deleted: int -# Called as chunk_writer(conn, repo_id, file_id, pf) once per file, inside the -# same conn.begin() as the rest of that file's row. Vectors must already be -# computed -- this seam never calls an embedder itself. -ChunkWriter = Callable[[Connection, int, int, ParsedFile], None] +# Called as chunk_writer(conn, repo_id, pairs) once per BATCH of files (a +# sequence of (file_id, pf) pairs), inside the same conn.begin() as the rest of +# those files' rows. Vectors must already be computed -- this seam never calls +# an embedder itself. +ChunkWriter = Callable[[Connection, int, Sequence[tuple[int, ParsedFile]]], None] # The two projection reads behind the file-level delta path. Returned as # ``(carried, present)`` -- see read_repo_content_shas. @@ -182,14 +214,19 @@ def index_repo( class is unioned in by ONE batched ``UPDATE ... RETURNING`` after the loop, which also supplies the ``file_id`` for its ``chunk_writer`` call (see :func:`_union_membership`). - * **changed/new** (everything else): an array-union upsert on - ``uq_files_repo_path_sha`` -- a file whose content already exists under - another branch gets THIS branch unioned into its ``branches`` array (one - row, shared content); a file whose content differs from every existing - version gets its own row. Then delete-and-reinsert its ``symbols`` and - ``reference_edges`` (neither has a natural key), then call - ``chunk_writer`` (if given) so chunk writes commit/roll back with the - rest of that file's row. + * **changed/new** (everything else): accumulated into an in-memory batch + and flushed (see :func:`_flush_file_batch`) once it reaches + ``_BATCH_MAX_FILES`` files or ``_BATCH_MAX_CONTENT_BYTES`` bytes, + whichever trips first, plus once more after the loop for whatever + remains. A flush issues ONE multi-row array-union upsert on + ``uq_files_repo_path_sha`` for the whole batch -- a file whose content + already exists under another branch gets THIS branch unioned into its + ``branches`` array (one row, shared content); a file whose content + differs from every existing version gets its own row -- then bulk + delete-and-reinsert of the batch's ``symbols`` and ``reference_edges`` + (neither has a natural key), then ONE ``chunk_writer`` call (if given) + for the whole batch, so every flushed file's rows commit/roll back + together as part of this same transaction. **Every** parsed file -- classified or written -- is collected into this branch's seen-set, so step 4 and its empty-seen-set guard are correct by @@ -319,6 +356,14 @@ class is unioned in by ONE batched ``UPDATE ... RETURNING`` after the # not the steady state, and chunk_writer's seam takes the ParsedFile. membership: list[tuple[ParsedFile, str]] = [] + # Changed/new files accumulate here and flush in batches of up to + # _BATCH_MAX_FILES files / _BATCH_MAX_CONTENT_BYTES bytes (whichever + # trips first), rather than issuing 5-7 statements per file. The check + # is POST-append, so a single file larger than the byte bound is never + # dropped or split -- it flushes with whatever batch it landed in. + batch: list[tuple[ParsedFile, FileExtraction, str]] = [] + batch_bytes = 0 + for pf, ex in items: sha = content_sha(pf.content) # Seen-set membership is recorded for EVERY parsed file, whatever its @@ -343,84 +388,33 @@ class is unioned in by ONE batched ``UPDATE ... RETURNING`` after the membership.append((pf, sha)) continue - file_stmt = ( - pg_insert(File) - .values( + batch.append((pf, ex, sha)) + batch_bytes += pf.size + if len(batch) >= _BATCH_MAX_FILES or batch_bytes >= _BATCH_MAX_CONTENT_BYTES: + s, e = _flush_file_batch( + conn, repo_id=repo_id, - path=pf.path, - lang=pf.lang, - size=pf.size, - content=pf.content, - commit=head_sha, - content_sha=sha, - branches=[branch], - ) - .on_conflict_do_update( - constraint="uq_files_repo_path_sha", - set_={ - "lang": pf.lang, - "size": pf.size, - "content": pf.content, - "commit": head_sha, - # Union this branch into whatever branches already share - # this exact content version -- a plain UNION via - # unnest+array_agg, row-lock-atomic regardless of - # concurrent readers (there is no concurrent WRITER for - # this repo -- see module docstring). - "branches": text( - "(SELECT array_agg(DISTINCT e) FROM " - "unnest(files.branches || excluded.branches) e)" - ), - }, - ) - .returning(File.id) - ) - file_id = conn.execute(file_stmt).scalar_one() - - conn.execute(delete(Symbol).where(Symbol.file_id == file_id)) - if ex.symbols: - conn.execute( - pg_insert(Symbol), - [ - { - "file_id": file_id, - "repo_id": repo_id, - "name": s.name, - "kind": s.kind, - "start_line": s.start_line, - "end_line": s.end_line, - } - for s in ex.symbols - ], - ) - symbol_count += len(ex.symbols) - - # UNCONDITIONAL, same as the symbols delete above: a file whose edges - # all vanish (e.g. every call/import site removed) must shed its stale - # rows even when this run's ex.edges is empty. - conn.execute(delete(ReferenceEdge).where(ReferenceEdge.file_id == file_id)) - if ex.edges: - conn.execute( - pg_insert(ReferenceEdge), - [ - { - "file_id": file_id, - "repo_id": repo_id, - "edge_kind": e.kind, - "target_name": e.target, - "line": e.line, - "enclosing_name": e.enclosing.name if e.enclosing else None, - "enclosing_kind": e.enclosing.kind if e.enclosing else None, - "enclosing_start_line": e.enclosing.start_line if e.enclosing else None, - "enclosing_end_line": e.enclosing.end_line if e.enclosing else None, - } - for e in ex.edges - ], + branch=branch, + head_sha=head_sha, + batch=batch, + chunk_writer=chunk_writer, ) - edge_count += len(ex.edges) + symbol_count += s + edge_count += e + batch = [] + batch_bytes = 0 - if chunk_writer is not None: - chunk_writer(conn, repo_id, file_id, pf) + if batch: + s, e = _flush_file_batch( + conn, + repo_id=repo_id, + branch=branch, + head_sha=head_sha, + batch=batch, + chunk_writer=chunk_writer, + ) + symbol_count += s + edge_count += e # ONE statement for the whole membership-only class, skipped entirely # when that class is empty (rather than issued as a no-op) so the @@ -486,6 +480,174 @@ class is unioned in by ONE batched ``UPDATE ... RETURNING`` after the return IndexCounts(files=file_count, symbols=symbol_count, swept=swept, edges=edge_count) +def _flush_file_batch( + conn: Connection, + *, + repo_id: int, + branch: str, + head_sha: str, + batch: list[tuple[ParsedFile, FileExtraction, str]], + chunk_writer: ChunkWriter | None, +) -> tuple[int, int]: + """Write one batch of changed/new files: one upsert, bulk delete+insert, one chunk_writer call. + + Returns ``(symbols_written, edges_written)`` to fold into ``index_repo``'s + running counts. Issues, in order: + + (a) **Intra-batch dedup guard**, keyed on ``(path, content_sha)``, keeping + the LAST occurrence (matching this module's existing per-file + last-write-wins) and logging one WARNING per collapsed key. Mandatory, + not defensive polish: a multi-row ``INSERT ... ON CONFLICT DO UPDATE`` + containing two rows with the same constrained values raises ``ON + CONFLICT DO UPDATE command cannot affect row a second time`` and + POISONS the transaction. A duplicate ``(path, content_sha)`` should be + impossible from the production source -- ``iter_tar_source_files`` + carries its own ``seen`` set and drops a repeat first-wins with a + WARNING (``indexer/ingest.py``) -- but ``items`` is an injected seam + other callers (tests, ``test_reconcile.py``) can feed directly. + (b) **One multi-row ``files`` upsert**, ``RETURNING id, path, + content_sha``. The ``SET`` clause uses ``excluded.*`` for every + per-row column (``lang``/``size``/``content``/``commit``) -- NEVER a + Python literal from one file, which would attach the LAST file's + values to every conflicting row in the batch: silent, committed, + corpus-wide corruption. ``commit`` becomes ``excluded.commit`` too + (every row in one batch shares ``head_sha`` so a literal would happen + to work here) so the rule stays uniform. ``branches`` keeps the same + array-union ``SET`` expression as the per-file upsert -- verified to + survive the multi-row form. + (c) **``DELETE ... WHERE file_id = ANY(:ids)``** for ``symbols`` and + ``reference_edges``, one statement each, over every id in this batch, + UNCONDITIONALLY (a file whose edges/symbols all vanished must still + shed its stale rows) -- then one param-budgeted bulk insert each via + :func:`indexer.bulk.insert_rows`. + (e) **One ``chunk_writer`` call** for the whole batch, given every + ``(file_id, pf)`` pair in batch order. + + **Ids are mapped from the upsert's ``RETURNING`` by ``(path, + content_sha)``, NEVER by row order** -- ``DO UPDATE ... RETURNING`` yields + one row per input row but makes no ordering guarantee. A missing key after + the map is built RAISES and rolls the whole ``(repo, branch)`` transaction + back (deliberately harsher than ``_union_membership``'s warn-and-skip: a + wrong or missing ``file_id`` here would attach one file's symbols to + another file's row -- durable core-corpus corruption, not a stale-vector + gap). + + Memory: this function's peak retention (``batch`` plus the row dicts built + below) is bounded by ``_BATCH_MAX_FILES`` / ``_BATCH_MAX_CONTENT_BYTES``, + ADDITIVE to ``indexer/ingest.py``'s ``MAX_FILE_BYTES`` + seen-path-set + retention (#106) -- this buffer is net-new retention, not a re-slicing of + memory the old extracted-tree path already held. + """ + # (a) Intra-batch dedup guard -- last occurrence wins. + deduped: dict[tuple[str, str], tuple[ParsedFile, FileExtraction, str]] = {} + for pf, ex, sha in batch: + key = (pf.path, sha) + if key in deduped: + logger.warning( + "duplicate (path, content_sha) %r within one batch; keeping the last occurrence", + key, + ) + deduped[key] = (pf, ex, sha) + entries = list(deduped.values()) + + # (b) One multi-row files upsert. excluded.* everywhere a per-file literal + # would otherwise leak the LAST file's values onto every conflicting row. + file_rows = [ + { + "repo_id": repo_id, + "path": pf.path, + "lang": pf.lang, + "size": pf.size, + "content": pf.content, + "commit": head_sha, + "content_sha": sha, + "branches": [branch], + } + for pf, _ex, sha in entries + ] + ins = pg_insert(cast(Table, File.__table__)).values(file_rows) + upsert_stmt = ins.on_conflict_do_update( + constraint="uq_files_repo_path_sha", + set_={ + "lang": ins.excluded.lang, + "size": ins.excluded.size, + "content": ins.excluded.content, + "commit": ins.excluded.commit, + "branches": text( + "(SELECT array_agg(DISTINCT e) FROM unnest(files.branches || excluded.branches) e)" + ), + }, + ).returning(File.id, File.path, File.content_sha) + returned = conn.execute(upsert_stmt).all() + file_ids = {(row.path, row.content_sha): row.id for row in returned} + + pairs: list[tuple[int, ParsedFile]] = [] + symbol_rows: list[dict[str, object]] = [] + edge_rows: list[dict[str, object]] = [] + ids: list[int] = [] + symbol_count = 0 + edge_count = 0 + for pf, ex, sha in entries: + file_id = file_ids.get((pf.path, sha)) + if file_id is None: + # NOT a warn-and-skip: an id missing from RETURNING means this + # file's symbols/edges could only be attached to the wrong row. + raise RuntimeError( + f"files upsert RETURNING has no row for (path={pf.path!r}, " + f"content_sha={sha!r}); refusing to attach its symbols/edges to another file" + ) + ids.append(file_id) + pairs.append((file_id, pf)) + symbol_rows.extend( + { + "file_id": file_id, + "repo_id": repo_id, + "name": s.name, + "kind": s.kind, + "start_line": s.start_line, + "end_line": s.end_line, + } + for s in ex.symbols + ) + symbol_count += len(ex.symbols) + edge_rows.extend( + { + "file_id": file_id, + "repo_id": repo_id, + "edge_kind": e.kind, + "target_name": e.target, + "line": e.line, + "enclosing_name": e.enclosing.name if e.enclosing else None, + "enclosing_kind": e.enclosing.kind if e.enclosing else None, + "enclosing_start_line": e.enclosing.start_line if e.enclosing else None, + "enclosing_end_line": e.enclosing.end_line if e.enclosing else None, + } + for e in ex.edges + ) + edge_count += len(ex.edges) + + # (c) Bulk delete-then-insert, unconditional (same semantics as the old + # per-file DELETE, which ran even for a file with zero symbols/edges). + conn.execute( + delete(Symbol).where(Symbol.file_id == any_(bindparam("ids", type_=ARRAY(BigInteger)))), + {"ids": ids}, + ) + insert_rows(conn, Symbol.__table__, symbol_rows) + + conn.execute( + delete(ReferenceEdge).where( + ReferenceEdge.file_id == any_(bindparam("ids", type_=ARRAY(BigInteger))) + ), + {"ids": ids}, + ) + insert_rows(conn, ReferenceEdge.__table__, edge_rows) + + if chunk_writer is not None: + chunk_writer(conn, repo_id, pairs) + + return symbol_count, edge_count + + def _repo_is_wholly_at_current_version(conn: Connection, *, repo_id: int) -> bool: """Statement 4: is EVERY ``repo_branches`` row for this repo at the current version? @@ -559,6 +721,7 @@ def _union_membership( if chunk_writer is None: return file_ids = {(row.path, row.content_sha): row.id for row in rows} + pairs: list[tuple[int, ParsedFile]] = [] for pf, sha in membership: file_id = file_ids.get((pf.path, sha)) if file_id is None: @@ -569,7 +732,9 @@ def _union_membership( pf.path, ) continue - chunk_writer(conn, repo_id, file_id, pf) + pairs.append((file_id, pf)) + if pairs: + chunk_writer(conn, repo_id, pairs) def _sweep_membership( diff --git a/tests/integration/AGENTS.md b/tests/integration/AGENTS.md index 85b6551..a3c053b 100644 --- a/tests/integration/AGENTS.md +++ b/tests/integration/AGENTS.md @@ -22,7 +22,9 @@ Within a run, isolation is per-test-module: every fixture creates a uniquely-nam | `test_semantic_rrf.py` | Hybrid RRF against the production `lakebase_ann`/`lakebase_bm25` operators: fusion plumbing (FULL OUTER JOIN + `1/(k+rank)`), real BM25 ranking, ANN index usage via EXPLAIN; `chunks` built with DDL identical to `0004`; fails loudly if the access methods are absent. | | `test_service.py` | `search_code_payload` keyset-cursor pagination: engine-per-call service, so the PGOPTIONS idiom (not a held-open connection) makes the schema visible to every pooled connection. | | `test_store.py` | `indexer.store.index_repo` upsert/sweep/rollback via the injected-connection seam; multi-branch scenarios: shared/divergent content across branches, per-branch CAS, empty-seen-set guard. | -| `test_store_chunk_writer.py` | `chunk_writer` seam end-to-end: chunks ride the same `conn.begin()` as the file row and cascade-delete on sweep (FK ON DELETE CASCADE); `chunks` table built with raw DDL (cross-`MetaData` FK can't be sorted by `create_all`). | +| `test_store_batching.py` | The batched changed/new write path (#105) against real Postgres: batch-size invariance ({1,2,7,500} produce a byte-identical corpus), the `excluded.*` trap behaviorally, row-dict column-completeness, mid-batch generator-failure rollback, measured statement-count-per-file (issue AC1), the multi-row branch-union upsert, and sweep/CAS still holding under batching. Same throwaway-schema idiom as `test_store_delta.py`, deliberately builds no `chunks` table. | +| `test_chunk_batching.py` | `indexer.chunk_store.write_chunks_batch` (#105) against real Postgres on the BARE `vector` extension (not `lakebase_vector`, unavailable locally) — one delete + one insert per batch, a zero-chunk covered file still deleted, an uncovered path untouched, FK cascade on sweep. Separate module from `test_store_batching.py` so a bare-`vector` fixture failure can't take that module's coverage down with it. | +| `test_store_chunk_writer.py` | `chunk_writer` seam end-to-end: chunks ride the same `conn.begin()` as the file row and cascade-delete on sweep (FK ON DELETE CASCADE); `chunks` table built with raw DDL (cross-`MetaData` FK can't be sorted by `create_all`). Lakebase-deferred (needs `lakebase_vector`); its stub's arity update to `(repo_id, pairs)` for #105 is reasoned through, not locally verified — see `test_chunk_batching.py` for the locally-runnable batched-chunk coverage. | | `test_symbols_search.py` | Symbol search (`sym:`): query → executed SQL → symbol defs; function-scoped `seeded` fixture (timeout + determinism tests need a clean corpus). | | `test_webui_semantic.py` | webui `/api/semantic` route on a Lakebase branch: dependency wiring, HTTP status/shape, not-migrated/disabled passthrough (ranking itself is `test_semantic_rrf.py`'s job); embedder seam monkeypatched, PGOPTIONS set before the overridden engine is built. | diff --git a/tests/integration/test_chunk_batching.py b/tests/integration/test_chunk_batching.py new file mode 100644 index 0000000..e2ba799 --- /dev/null +++ b/tests/integration/test_chunk_batching.py @@ -0,0 +1,298 @@ +"""Integration tests for indexer.chunk_store.write_chunks_batch (#105) against real Postgres. + +**Deliberately builds ``chunks`` on the BARE ``vector`` extension**, not +``lakebase_vector``: `codesearch-pg` (this repo's local dev Postgres) has +``vector 0.8.5`` installed, and ``lakebase_vector`` depends on that same base +extension (``app/alembic/versions/0004_semantic_chunks.py``). What actually +fails locally is ``CREATE EXTENSION IF NOT EXISTS lakebase_vector CASCADE`` +(the beta extension itself, unavailable outside Lakebase) -- +``tests/integration/test_store_chunk_writer.py``'s whole module is +Lakebase-deferred for exactly that reason. This module's DDL diverges from +production in two ways -- a plain nullable ``ts tsvector`` (not a +``GENERATED`` column) and no ANN index -- and that divergence is IRRELEVANT to +what this module proves: ``write_chunks_batch``'s delete-by-``ANY`` and +multi-row insert semantics depend on the ``embedding`` column's TYPE, not on +whether an index exists over it or on how ``ts`` is populated. A separate +module from ``tests/integration/test_store_batching.py`` on purpose, so a +failure in this bare-``vector`` fixture can never take that module's coverage +down with it. +""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence + +import pytest +from sqlalchemy import Connection, text + +from app.config import SEMANTIC_EMBEDDING_DIM +from app.db.client import create_db_engine +from app.db.models import Base +from indexer.chunk_store import write_chunks_batch +from indexer.languages import ExtractedSymbol, FileExtraction, ParsedFile +from indexer.store import index_repo + +SCHEMA = "test_chunk_batching" + + +@pytest.fixture +def conn() -> Iterator[Connection]: + engine = create_db_engine() + connection = engine.connect() + try: + connection.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm")) + connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + connection.execute(text(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE")) + connection.execute(text(f"CREATE SCHEMA {SCHEMA}")) + connection.execute(text(f"SET search_path TO {SCHEMA}, public")) + connection.commit() + + Base.metadata.create_all(bind=connection) + connection.execute( + text( + "CREATE TABLE chunks (" + "id bigserial PRIMARY KEY, " + "file_id integer NOT NULL REFERENCES files(id) ON DELETE CASCADE, " + "chunk_index integer NOT NULL, " + "content text NOT NULL, " + "start_line integer, " + "end_line integer, " + f"embedding vector({SEMANTIC_EMBEDDING_DIM}), " + "ts tsvector, " + "CONSTRAINT uq_chunks_file_id_chunk_index UNIQUE (file_id, chunk_index))" + ) + ) + connection.commit() + + yield connection + finally: + connection.rollback() + connection.execute(text(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE")) + connection.commit() + connection.close() + engine.dispose() + + +_STUB_VECTOR = [0.1] * SEMANTIC_EMBEDDING_DIM + + +def _seed_repo(conn: Connection, name: str = "acme/widgets") -> int: + return int( + conn.execute( + text("INSERT INTO repos (name) VALUES (:name) RETURNING id"), {"name": name} + ).scalar_one() + ) + + +def _seed_file(conn: Connection, repo_id: int, path: str) -> int: + return int( + conn.execute( + text( + "INSERT INTO files " + "(repo_id, path, lang, size, content, commit, content_sha, branches) " + "VALUES (:repo_id, :path, 'python', 1, 'x', 'sha', :sha, ARRAY['main']) " + "RETURNING id" + ), + {"repo_id": repo_id, "path": path, "sha": path}, + ).scalar_one() + ) + + +def _chunk_count(conn: Connection, file_id: int) -> int: + return int( + conn.execute( + text("SELECT count(*) FROM chunks WHERE file_id = :id"), {"id": file_id} + ).scalar_one() + ) + + +def _pf(path: str, content: str) -> ParsedFile: + return ParsedFile(path=path, lang="python", size=len(content.encode()), content=content) + + +def _items( + *specs: tuple[str, str, list[ExtractedSymbol]], +) -> list[tuple[ParsedFile, FileExtraction]]: + return [ + (_pf(path, content), FileExtraction(symbols=syms, edges=[])) + for path, content, syms in specs + ] + + +MAIN = ("main.py", "def f():\n return 1\n", [ExtractedSymbol("f", "function", 1, 2)]) +UTIL = ("util.py", "def g():\n return 2\n", [ExtractedSymbol("g", "function", 1, 2)]) + + +def _stub_chunk_writer( + conn: Connection, repo_id: int, pairs: Sequence[tuple[int, ParsedFile]] +) -> None: + write_chunks_batch( + conn, + rows=[(file_id, [(0, pf.content, 1, 2, _STUB_VECTOR)]) for file_id, pf in pairs], + ) + + +# --- Test 20: one delete + one insert per batch; zero-chunk covered files ---- +# still deleted; uncovered paths untouched ----------------------------------- + + +@pytest.mark.integration +def test_batch_write_deletes_zero_chunk_covered_files_and_leaves_uncovered_untouched( + conn: Connection, +) -> None: + repo_id = _seed_repo(conn) + covered_a = _seed_file(conn, repo_id, "a.py") + covered_b = _seed_file(conn, repo_id, "b.py") # covered this run, but zero NEW chunks + uncovered = _seed_file(conn, repo_id, "c.py") # not part of this batch at all + conn.commit() + + # Seed b.py and c.py with a prior chunk row each. + write_chunks_batch(conn, rows=[(covered_b, [(0, "old", 1, 1, _STUB_VECTOR)])]) + write_chunks_batch(conn, rows=[(uncovered, [(0, "old", 1, 1, _STUB_VECTOR)])]) + conn.commit() + + written = write_chunks_batch( + conn, + rows=[ + (covered_a, [(0, "a chunk", 1, 1, _STUB_VECTOR)]), + (covered_b, []), + ], + ) + assert written == 1 + + assert _chunk_count(conn, covered_a) == 1 + assert _chunk_count(conn, covered_b) == 0 # deleted, not left stale + assert _chunk_count(conn, uncovered) == 1 # untouched -- not in this batch + + +# --- Test 21: FK cascade still removes chunks when the sweep deletes a file - + + +# --- membership-only + real chunks: the case tests/integration/test_store_chunk_writer.py +# cannot exercise locally (Lakebase-deferred), covered here since this module's +# bare-`vector` fixture actually runs against codesearch-pg. ----------------- + + +@pytest.mark.integration +def test_membership_only_acquisition_writes_chunks_for_the_acquiring_branch( + conn: Connection, +) -> None: + """Branch 'b' acquires MAIN's content already stored under branch 'a' via + the membership-only path (``indexer.store._union_membership``) -- no + symbol/edge rewrite, but chunks ARE written for the acquiring branch via + the same ``write_chunks_batch`` this module otherwise tests. + """ + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a1", + items=_items(MAIN), + chunk_writer=_stub_chunk_writer, + ) + conn.rollback() + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b1", + items=_items(UTIL), + chunk_writer=_stub_chunk_writer, + ) + conn.rollback() + + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b2", + items=_items(MAIN, UTIL), + chunk_writer=_stub_chunk_writer, + ) + main_file_id = conn.execute(text("SELECT id FROM files WHERE path = 'main.py'")).scalar_one() + assert _chunk_count(conn, main_file_id) == 1 + + +@pytest.mark.integration +def test_membership_only_duplicate_item_does_not_poison_the_chunk_insert( + conn: Connection, +) -> None: + """A duplicated ``(path, content_sha)`` entry in ``items`` landing in the + membership-only class must not raise a UNIQUE VIOLATION against + ``uq_chunks_file_id_chunk_index``. ``write_chunks_batch``'s dedup guard + (mirroring ``indexer.store._flush_file_batch``'s own guard) is what + prevents it -- without it, ``_union_membership`` would hand the same + ``file_id`` to ``chunk_writer`` twice, and the batch insert would attempt + two rows with the same ``(file_id, chunk_index)``. + + ``items`` is an injected seam (this module, ``tests/integration/test_reconcile.py``, + and any other direct caller can feed it a duplicate) even though the + production source (``indexer/ingest.py``'s ``iter_tar_source_files``) never + does. + """ + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a1", + items=_items(MAIN), + chunk_writer=_stub_chunk_writer, + ) + conn.rollback() + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b1", + items=_items(UTIL), + chunk_writer=_stub_chunk_writer, + ) + conn.rollback() + + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b2", + items=_items(MAIN, MAIN, UTIL), + chunk_writer=_stub_chunk_writer, + ) + main_file_id = conn.execute(text("SELECT id FROM files WHERE path = 'main.py'")).scalar_one() + assert _chunk_count(conn, main_file_id) == 1 + + +@pytest.mark.integration +def test_fk_cascade_removes_chunk_rows_when_sweep_deletes_an_emptied_file( + conn: Connection, +) -> None: + index_repo( + conn, + name="acme/widgets", + branch="main", + is_default=True, + head_sha="sha1", + items=_items(MAIN, UTIL), + chunk_writer=_stub_chunk_writer, + ) + util_file_id = conn.execute(text("SELECT id FROM files WHERE path = 'util.py'")).scalar_one() + assert _chunk_count(conn, util_file_id) == 1 + conn.rollback() + + counts = index_repo( + conn, + name="acme/widgets", + branch="main", + is_default=True, + head_sha="sha2", + items=_items(MAIN), # util.py dropped -> swept + chunk_writer=_stub_chunk_writer, + ) + assert counts.swept == 1 + assert conn.execute(text("SELECT count(*) FROM files WHERE path = 'util.py'")).scalar_one() == 0 + assert _chunk_count(conn, util_file_id) == 0 # FK ON DELETE CASCADE diff --git a/tests/integration/test_reconcile.py b/tests/integration/test_reconcile.py index 4e98cbd..3abbdd7 100644 --- a/tests/integration/test_reconcile.py +++ b/tests/integration/test_reconcile.py @@ -13,7 +13,7 @@ import os import uuid -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from typing import Any import pytest @@ -24,7 +24,7 @@ from app.db.client import create_db_engine from app.db.grants import build_job_grants from app.db.models import Base -from indexer.chunk_store import write_chunks +from indexer.chunk_store import write_chunks_batch from indexer.languages import ExtractedSymbol, FileExtraction, ParsedFile from indexer.store import ( ReconcileCounts, @@ -98,8 +98,18 @@ def _items( _STUB_VECTOR = [0.1] * SEMANTIC_EMBEDDING_DIM -def _stub_chunk_writer(conn: Connection, repo_id: int, file_id: int, pf: ParsedFile) -> None: - write_chunks(conn, file_id=file_id, chunks=[(0, pf.content, 1, 2, _STUB_VECTOR)]) +def _stub_chunk_writer( + conn: Connection, repo_id: int, pairs: Sequence[tuple[int, ParsedFile]] +) -> None: + # One write_chunks_batch call for the whole batch, matching indexer/job.py's + # real closure and #105's reshaped per-BATCH seam -- Lakebase-deferred, this + # module cannot run locally (needs lakebase_vector), so this update is + # reasoned through against indexer/store.py's ChunkWriter alias, never + # locally verified. + write_chunks_batch( + conn, + rows=[(file_id, [(0, pf.content, 1, 2, _STUB_VECTOR)]) for file_id, pf in pairs], + ) def _count(conn: Connection, table: str, where: str = "") -> int: diff --git a/tests/integration/test_store_batching.py b/tests/integration/test_store_batching.py new file mode 100644 index 0000000..7128906 --- /dev/null +++ b/tests/integration/test_store_batching.py @@ -0,0 +1,490 @@ +"""Integration tests for the batched changed/new write path (#105) against real Postgres. + +Row-identity and byte-parity proof against real SQL, complementing +``tests/unit/test_store_batching.py``'s statement-level pins. Clones +``tests/integration/test_store_delta.py``'s throwaway-schema idiom (own copy, +per ``tests/integration/AGENTS.md``'s no-conftest convention). + +**Deliberately builds no ``chunks`` table and creates no ``lakebase_*`` +extension** -- same discipline as ``test_store_delta.py``, for the same reason +(fixtures are module-local by convention, so one failing fixture would +vaporize the whole module). Chunk-touching batched cases live in +``tests/integration/test_chunk_batching.py`` instead. +""" + +from __future__ import annotations + +import contextlib +import re +from collections.abc import Iterator +from typing import Any +from uuid import uuid4 + +import pytest +from sqlalchemy import Connection, event, text + +import indexer.store as store_module +from app.db.client import create_db_engine +from app.db.models import INDEX_SEMANTICS_VERSION, Base, File +from indexer.languages import ( + ExtractedEdge, + ExtractedSymbol, + FileExtraction, + IndexCounts, + ParsedFile, +) +from indexer.store import StaleIndexError, _stamp_repo_branch, index_repo + +SCHEMA_PREFIX = "test_store_batching" + + +@pytest.fixture +def conn() -> Iterator[Connection]: + with _fresh_schema() as connection: + yield connection + + +@contextlib.contextmanager +def _fresh_schema() -> Iterator[Connection]: + """A uniquely-named throwaway schema + a live connection, torn down after.""" + schema = f"{SCHEMA_PREFIX}_{uuid4().hex[:12]}" + engine = create_db_engine() + connection = engine.connect() + try: + connection.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + connection.execute(text(f"CREATE SCHEMA {schema}")) + connection.execute(text(f"SET search_path TO {schema}, public")) + connection.commit() + + Base.metadata.create_all(bind=connection) + connection.commit() + + yield connection + finally: + connection.rollback() + connection.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE")) + connection.commit() + connection.close() + engine.dispose() + + +def _pf(path: str, content: str) -> ParsedFile: + return ParsedFile(path=path, lang="python", size=len(content.encode()), content=content) + + +def _items( + *specs: tuple[str, str, list[ExtractedSymbol]], +) -> list[tuple[ParsedFile, FileExtraction]]: + return [ + (_pf(path, content), FileExtraction(symbols=syms, edges=[])) + for path, content, syms in specs + ] + + +def _fn(name: str, n: int) -> tuple[str, str, list[ExtractedSymbol]]: + content = f"def f{n}():\n return {n}\n" + return (f"{name}{n}.py", content, [ExtractedSymbol(f"f{n}", "function", 1, 2)]) + + +def _index_default( + conn: Connection, *, name: str, head_sha: str, items: list[tuple[ParsedFile, FileExtraction]] +) -> IndexCounts: + return index_repo( + conn, name=name, branch="main", is_default=True, head_sha=head_sha, items=items + ) + + +def _count(conn: Connection, table: str, where: str = "") -> int: + sql = f"SELECT count(*) FROM {table}" + if where: + sql += f" WHERE {where}" + return int(conn.execute(text(sql)).scalar_one()) + + +def _sym(prefix: str, n: int) -> ExtractedSymbol: + return ExtractedSymbol(f"{prefix}{n}", "function", 1, 2) + + +def _fixture_items(branch: str) -> list[tuple[ParsedFile, FileExtraction]]: + """~15 files: 5 content-deduped across every branch, 2 divergent per branch, + one zero-symbols file, one file with a real edge, and 6 branch-unique files + -- enough diversity to exercise every row shape the batched write touches.""" + items: list[tuple[ParsedFile, FileExtraction]] = [] + + for i in range(5): + content = f"def shared{i}():\n return {i}\n" + extraction = FileExtraction(symbols=[_sym("shared", i)], edges=[]) + items.append((_pf(f"shared{i}.py", content), extraction)) + + for i in range(2): + # A deterministic per-(branch, i) value -- NOT Python's hash(), which is + # randomized per-process (PYTHONHASHSEED) and would make this fixture's + # content non-reproducible across separate runs/processes. + magic = (sum(ord(c) for c in branch) * 31 + i) % 997 + content = f"def divergent{i}():\n return {magic}\n" + extraction = FileExtraction(symbols=[_sym("divergent", i)], edges=[]) + items.append((_pf(f"divergent{i}.py", content), extraction)) + + items.append((_pf("nosymbols.py", "# just a comment\n"), FileExtraction(symbols=[], edges=[]))) + + caller = _sym("caller", 0) + items.append( + ( + _pf(f"{branch}_withedges.py", "def caller0():\n callee()\n"), + FileExtraction( + symbols=[caller], + edges=[ExtractedEdge(kind="call", target="callee", line=2, enclosing=caller)], + ), + ) + ) + + for i in range(6): + content = f"def {branch}_only{i}():\n return {i}\n" + extraction = FileExtraction(symbols=[_sym(f"{branch}_only", i)], edges=[]) + items.append((_pf(f"{branch}_only{i}.py", content), extraction)) + + return items + + +def _dump_files(conn: Connection) -> list[tuple[Any, ...]]: + rows = conn.execute( + text( + "SELECT repo_id, path, content_sha, lang, size, content, commit, branches " + "FROM files ORDER BY path, content_sha" + ) + ).all() + return sorted( + ( + r.repo_id, + r.path, + r.content_sha, + r.lang, + r.size, + r.content, + r.commit, + tuple(sorted(r.branches)), + ) + for r in rows + ) + + +def _dump_symbols(conn: Connection) -> list[tuple[Any, ...]]: + rows = conn.execute( + text( + "SELECT f.path, s.name, s.kind, s.start_line, s.end_line " + "FROM symbols s JOIN files f ON f.id = s.file_id" + ) + ).all() + return sorted(tuple(r) for r in rows) + + +def _dump_edges(conn: Connection) -> list[tuple[Any, ...]]: + rows = conn.execute( + text( + "SELECT f.path, e.edge_kind, e.target_name, e.line, " + "e.enclosing_name, e.enclosing_kind, " + "e.enclosing_start_line, e.enclosing_end_line " + "FROM reference_edges e JOIN files f ON f.id = e.file_id" + ) + ).all() + return sorted(tuple(r) for r in rows) + + +# --- Test 14: batch-size invariance is the parity harness -------------------- + + +@pytest.mark.integration +def test_batch_size_invariance_produces_a_byte_identical_corpus( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Index the same 3-branch fixture repo at _BATCH_MAX_FILES in {1, 2, 7, 500} + into four throwaway schemas; the normalized (files, symbols, edges) dump + must be identical across all four -- the direct proof that batch size + changes round trips, never corpus content.""" + dumps = [] + for batch_size in (1, 2, 7, 500): + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", batch_size) + with _fresh_schema() as conn: + index_repo( + conn, + name="acme/widgets", + branch="a", + is_default=True, + head_sha="sha_a", + items=_fixture_items("a"), + ) + conn.commit() + index_repo( + conn, + name="acme/widgets", + branch="b", + is_default=False, + head_sha="sha_b", + items=_fixture_items("b"), + ) + conn.commit() + index_repo( + conn, + name="acme/widgets", + branch="c", + is_default=False, + head_sha="sha_c", + items=_fixture_items("c"), + ) + conn.commit() + dumps.append((_dump_files(conn), _dump_symbols(conn), _dump_edges(conn))) + + first = dumps[0] + for batch_size, dump in zip((2, 7, 500), dumps[1:], strict=True): + assert dump == first, f"corpus at _BATCH_MAX_FILES={batch_size} diverged from size=1" + + +# --- Test 15: the excluded.* trap, behaviorally ------------------------------- + + +@pytest.mark.integration +def test_excluded_trap_every_conflicting_row_keeps_its_own_values(conn: Connection) -> None: + """Re-index a batch of 3 files whose contents ALL differ, where every row + already exists (every row takes DO UPDATE): each row's stored + content/lang/size/commit must be its OWN, not the last file's. Under the + literal-binding bug all three collapse to the last file's values.""" + v1 = _items(("a.py", "x = 1\n", []), ("b.py", "y = 2\n", []), ("c.py", "z = 3\n", [])) + _index_default(conn, name="acme/widgets", head_sha="sha1", items=v1) + conn.commit() + + v2 = _items(("a.py", "x = 100\n", []), ("b.py", "y = 200\n", []), ("c.py", "z = 300\n", [])) + _index_default(conn, name="acme/widgets", head_sha="sha2", items=v2) + + for path, content in [("a.py", "x = 100\n"), ("b.py", "y = 200\n"), ("c.py", "z = 300\n")]: + row = conn.execute( + text("SELECT content, commit, size FROM files WHERE path = :p"), {"p": path} + ).one() + assert row.content == content + assert row.commit == "sha2" + assert row.size == len(content.encode()) + + +# --- Test 15b: column-completeness of the batched upsert's row dict --------- + + +@pytest.mark.integration +def test_files_upsert_row_dict_covers_every_non_id_column(conn: Connection) -> None: + """The row-dict key set built by _flush_file_batch must equal every `files` + column except `id` -- catches a column silently dropped uniformly at every + batch size, which test 14's batched-against-batched comparison cannot.""" + captured_sql: list[str] = [] + + def _listener( + conn_: Any, cursor: Any, statement: str, parameters: Any, context: Any, executemany: bool + ) -> None: + if statement.startswith("INSERT INTO files"): + captured_sql.append(statement) + + event.listen(conn.engine, "before_cursor_execute", _listener) + try: + _index_default( + conn, + name="acme/widgets", + head_sha="sha1", + items=_items(("a.py", "x = 1\n", []), ("b.py", "y = 2\n", [])), + ) + finally: + event.remove(conn.engine, "before_cursor_execute", _listener) + + assert captured_sql + match = re.search(r"INSERT INTO files \(([^)]+)\)", captured_sql[0]) + assert match is not None + columns = {c.strip() for c in match.group(1).split(",")} + assert columns == set(File.__table__.columns.keys()) - {"id"} + + +# --- Test 16: mid-batch rollback ---------------------------------------------- + + +@pytest.mark.integration +def test_mid_batch_generator_failure_rolls_back_the_whole_transaction( + conn: Connection, monkeypatch: pytest.MonkeyPatch +) -> None: + """An `items` generator that raises AFTER at least one flush already ran + (_BATCH_MAX_FILES=5, 6 files yielded before the raise) leaves ZERO + files/symbols/reference_edges rows committed and no repo_branches row -- + the whole (repo, branch) transaction rolls back, not just the failing batch.""" + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", 5) + + class _Boom(Exception): + pass + + def _raising_items() -> Iterator[tuple[ParsedFile, FileExtraction]]: + for i in range(6): + yield _pf(f"f{i}.py", f"x = {i}\n"), FileExtraction(symbols=[], edges=[]) + raise _Boom("simulated failure mid-generator, after the first flush") + + with pytest.raises(_Boom): + index_repo( + conn, + name="acme/widgets", + branch="main", + is_default=True, + head_sha="sha1", + items=_raising_items(), + ) + + assert _count(conn, "files") == 0 + assert _count(conn, "symbols") == 0 + assert _count(conn, "reference_edges") == 0 + assert _count(conn, "repo_branches") == 0 + + +# --- Test 17: statement count, measured --------------------------------------- + + +@pytest.mark.integration +def test_statement_count_per_file_meets_the_acceptance_criterion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """AC 1: round trips per file drop from 3-7 to <= 0.05. Measures real cursor + executions over a 600-file first-time index, both batched and at + _BATCH_MAX_FILES=1, and prints the measured value for the PR body.""" + n = 600 + fixture_items = _items(*[_fn("f", i) for i in range(n)]) + + def _count_statements(batch_size: int) -> int: + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", batch_size) + with _fresh_schema() as conn: + count = 0 + + def _listener( + conn_: Any, + cursor: Any, + statement: str, + parameters: Any, + context: Any, + executemany: bool, + ) -> None: + nonlocal count + count += 1 + + event.listen(conn.engine, "before_cursor_execute", _listener) + try: + index_repo( + conn, + name="acme/widgets", + branch="main", + is_default=True, + head_sha="sha1", + items=fixture_items, + ) + finally: + event.remove(conn.engine, "before_cursor_execute", _listener) + return count + + batched = _count_statements(500) + unbatched = _count_statements(1) + + # Fixed per-transaction overhead independent of N and of batch size: + # repos-insert, repo_branches-insert, sweep-update, sweep-delete, stamp. + fixed_overhead = 5 + per_file = (batched - fixed_overhead) / n + print(f"\nstatements/file at _BATCH_MAX_FILES=500, N={n}: {per_file:.4f} (issue AC: <= 0.05)") + assert per_file <= 0.05 + assert batched < unbatched + + +# --- Test 18: branch union survives the multi-row form, gate closed --------- + + +@pytest.mark.integration +def test_branch_union_survives_the_multi_row_form_with_the_gate_closed(conn: Connection) -> None: + """Branch 'b' indexes a batch of files whose exact (path, content) already + exist under branch 'a'. 'b' has never indexed before (baseline_version is + NULL), so the delta gate is CLOSED and every file takes the changed/new + batched path -- not membership-only. Every such row's `branches` must end + ['a', 'b'], sorted-distinct: the multi-row array-union SET, not a + per-row overwrite.""" + shared = _items(*[_fn("shared", i) for i in range(5)]) + index_repo( + conn, name="acme/widgets", branch="a", is_default=True, head_sha="sha_a", items=shared + ) + conn.commit() + + index_repo( + conn, name="acme/widgets", branch="b", is_default=False, head_sha="sha_b", items=shared + ) + + rows = conn.execute(text("SELECT path, branches FROM files")).all() + assert len(rows) == 5 + for row in rows: + assert sorted(row.branches) == ["a", "b"] + + +# --- Test 19: sweep and CAS still hold under batching ------------------------ + + +@pytest.mark.integration +def test_sweep_still_removes_a_deleted_file_under_batching( + conn: Connection, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", 2) + items_v1 = _items(*[_fn("f", i) for i in range(5)]) + index_repo( + conn, name="acme/widgets", branch="main", is_default=True, head_sha="sha1", items=items_v1 + ) + conn.commit() + + items_v2 = _items(*[_fn("f", i) for i in range(4)]) # f4.py dropped + counts = index_repo( + conn, name="acme/widgets", branch="main", is_default=True, head_sha="sha2", items=items_v2 + ) + assert counts.swept == 1 + assert _count(conn, "files", "path = 'f4.py'") == 0 + + +@pytest.mark.integration +def test_cas_conflict_rolls_back_the_whole_batched_transaction( + conn: Connection, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", 2) + items = _items(*[_fn("f", i) for i in range(5)]) + index_repo( + conn, name="acme/widgets", branch="main", is_default=True, head_sha="sha1", items=items + ) + files_before = _count(conn, "files") + repo_id = int( + conn.execute(text("SELECT id FROM repos WHERE name = 'acme/widgets'")).scalar_one() + ) + conn.rollback() + + with pytest.raises(StaleIndexError, match="wrong_sha"), conn.begin(): + conn.execute(text("DELETE FROM files WHERE path = 'f0.py'")) + _stamp_repo_branch( + conn, + name="acme/widgets", + branch="main", + repo_id=repo_id, + head_sha="sha2", + baseline_commit="wrong_sha", + baseline_version=INDEX_SEMANTICS_VERSION, + ) + assert _count(conn, "files") == files_before + + +@pytest.mark.integration +def test_empty_items_skips_sweep_and_holds_the_semantics_version_under_batching( + conn: Connection, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", 2) + items = _items(*[_fn("f", i) for i in range(5)]) + index_repo( + conn, name="acme/widgets", branch="main", is_default=True, head_sha="sha1", items=items + ) + conn.execute(text("UPDATE repo_branches SET index_semantics_version = 3")) + conn.commit() + + counts = index_repo( + conn, name="acme/widgets", branch="main", is_default=True, head_sha="sha2", items=[] + ) + assert counts == IndexCounts(files=0, symbols=0, swept=0, edges=0) + stamp = conn.execute( + text("SELECT last_indexed_commit, index_semantics_version FROM repo_branches") + ).one() + assert stamp == ("sha2", 3) diff --git a/tests/integration/test_store_chunk_writer.py b/tests/integration/test_store_chunk_writer.py index 30815b4..6a8308d 100644 --- a/tests/integration/test_store_chunk_writer.py +++ b/tests/integration/test_store_chunk_writer.py @@ -27,7 +27,7 @@ from __future__ import annotations -from collections.abc import Iterator +from collections.abc import Iterator, Sequence import pytest from sqlalchemy import Connection, text @@ -35,7 +35,7 @@ from app.config import SEMANTIC_EMBEDDING_DIM from app.db.client import create_db_engine from app.db.models import Base -from indexer.chunk_store import write_chunks +from indexer.chunk_store import write_chunks_batch from indexer.languages import ExtractedSymbol, FileExtraction, IndexCounts, ParsedFile from indexer.store import index_repo @@ -88,10 +88,20 @@ def _pf(path: str, content: str) -> ParsedFile: _STUB_VECTOR = [0.1] * SEMANTIC_EMBEDDING_DIM -def _stub_chunk_writer(conn: Connection, repo_id: int, file_id: int, pf: ParsedFile) -> None: +def _stub_chunk_writer( + conn: Connection, repo_id: int, pairs: Sequence[tuple[int, ParsedFile]] +) -> None: # A fixed, precomputed 1-chunk-per-file "embedding" -- proves the seam without - # needing a real embedder (chunk_writer never calls one). - write_chunks(conn, file_id=file_id, chunks=[(0, pf.content, 1, 2, _STUB_VECTOR)]) + # needing a real embedder (chunk_writer never calls one). One write_chunks_batch + # call for the whole batch, matching indexer/job.py's real closure and #105's + # reshaped per-BATCH seam (tests/integration/test_chunk_batching.py exercises + # write_chunks_batch directly; this module still can't run locally -- + # lakebase_vector, see the module docstring -- so this only fixes what a + # future Lakebase run would exercise, it does not itself verify anything here). + write_chunks_batch( + conn, + rows=[(file_id, [(0, pf.content, 1, 2, _STUB_VECTOR)]) for file_id, pf in pairs], + ) def _count(conn: Connection, table: str, where: str = "") -> int: @@ -234,10 +244,11 @@ def test_unchanged_file_never_calls_chunk_writer_and_preserves_chunk_ids( calls: list[str] = [] def _tracking_chunk_writer( - conn: Connection, repo_id: int, file_id: int, pf: ParsedFile + conn: Connection, repo_id: int, pairs: Sequence[tuple[int, ParsedFile]] ) -> None: - calls.append(pf.path) - _stub_chunk_writer(conn, repo_id, file_id, pf) + for file_id, pf in pairs: + calls.append(pf.path) + _stub_chunk_writer(conn, repo_id, [(file_id, pf)]) index_repo( conn, diff --git a/tests/unit/AGENTS.md b/tests/unit/AGENTS.md index c12bc22..57151a7 100644 --- a/tests/unit/AGENTS.md +++ b/tests/unit/AGENTS.md @@ -11,7 +11,8 @@ Hermetic unit tests: no network, no database, no Databricks SDK instantiation. E |------|-------------| | `__init__.py` | Empty package marker. | | `test_branches.py` | `indexer.branches.resolve_branches`: glob matching, dedup, cap; empty globs → default branch only. | -| `test_chunk_store.py` | `indexer.chunk_store.write_chunks` statement shape via a fake `Connection` recording `execute` calls; delete-then-insert, and proof no embedding call happens here. | +| `test_bulk.py` | `indexer.bulk.insert_rows`: param-budgeted slicing into multi-row `VALUES` statements, no DB required. | +| `test_chunk_store.py` | `indexer.chunk_store.write_chunks`/`write_chunks_batch` statement shape via a fake `Connection` recording `execute` calls; delete-then-insert (one `DELETE ... WHERE file_id = ANY(:ids)` since #105), and proof no embedding call happens here. | | `test_chunking.py` | `indexer.parse.iter_chunks` chunking behavior. | | `test_ci_branch.py` | `scripts/ci_branch.py` lifecycle with the SDK fully faked; pins that teardown NEVER raises and every create carries a TTL (leak protection for cancelled CI runs). | | `test_db_client.py` | Engine factory local (`PGHOST`) mode builds without instantiating the SDK; ORM models expose exactly the durable-core columns, constraints, and GIN indexes. | @@ -38,7 +39,8 @@ Hermetic unit tests: no network, no database, no Databricks SDK instantiation. E | `test_semantics_version_tripwire.py` | CI tripwire: extraction-semantics changes must bump `INDEX_SEMANTICS_VERSION`. | | `test_service.py` | Keyset-cursor pagination in `app/service.py`: pure cursor encode/decode + pagination-mode gating with fake engine/`GrepResult` (real keyset SQL lives in integration). | | `test_smoke.py` | Pure predicate functions in `scripts/smoke.py`, loaded by file path (scripts/ is not a package). | -| `test_store_chunk_writer.py` | `indexer.store`'s optional `chunk_writer` param with a hand-rolled fake `Connection`: writer called inside the same `conn.begin()` with `(repo_id, file_id, pf)`. | +| `test_store_batching.py` | `indexer.store`'s batched changed/new write path (#105) at the statement level: boundary tripping by count and by bytes, the `excluded.*` trap, id mapping by `(path, content_sha)` (never row order), the intra-batch dedup guard, and the exact statement inventory of a mixed run. Extends `test_store_delta.py`'s `_FakeConn` idiom with a `files_returning`/`membership_returning` hook to script `RETURNING` row order independently of insertion order. | +| `test_store_chunk_writer.py` | `indexer.store`'s optional `chunk_writer` param with a hand-rolled fake `Connection`: writer called inside the same `conn.begin()` with `(repo_id, pairs)`, `pairs` a sequence of `(file_id, pf)` (one call per batch since #105). | | `test_symbols.py` | `indexer.symbols.extract_symbols` across the V1 languages; nested symbols and the Python-vs-JS/TS method-kind asymmetry. | | `test_symbols_search.py` | `sym:` atom walker (pure) + rendered step-2 projection SQL; the composed two-query path is integration's job. | | `test_webui_main.py` | Route-level tests of the webui FastAPI backend via `app.dependency_overrides` of the `get_engine`/`get_settings` dependencies. | diff --git a/tests/unit/test_bulk.py b/tests/unit/test_bulk.py new file mode 100644 index 0000000..c01fd77 --- /dev/null +++ b/tests/unit/test_bulk.py @@ -0,0 +1,86 @@ +"""Unit tests for indexer.bulk.insert_rows: statement shape and slicing, no DB required. + +A fake ``Connection`` records the compiled ``Insert`` construct passed to +``execute`` so the multi-row-VALUES shape (as opposed to an executemany param +list) and the exact row count per statement can be asserted without a real +Postgres. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from sqlalchemy import Column, Integer, MetaData, Table +from sqlalchemy.dialects.postgresql import dialect as pg_dialect + +from indexer.bulk import insert_rows + +_METADATA = MetaData() +_TABLE = Table("widgets", _METADATA, Column("a", Integer), Column("b", Integer)) + + +class _FakeConn: + def __init__(self) -> None: + self.statements: list[Any] = [] + + def execute(self, stmt: Any, params: Any = None) -> None: + assert params is None # rows travel inside stmt.values(...), not as a param list + self.statements.append(stmt) + + +def _rows(n: int) -> list[dict[str, int]]: + return [{"a": i, "b": i * 10} for i in range(n)] + + +@pytest.mark.unit +def test_empty_list_issues_zero_statements() -> None: + conn = _FakeConn() + insert_rows(conn, _TABLE, []) + assert conn.statements == [] + + +@pytest.mark.unit +def test_slicing_produces_exact_ceil_statements_every_row_once_in_order() -> None: + # 2 columns, budget 10 -> 5 rows/statement. 12 rows -> ceil(12/5) = 3 statements + # of sizes 5, 5, 2. + conn = _FakeConn() + rows = _rows(12) + insert_rows(conn, _TABLE, rows, param_budget=10) + + assert len(conn.statements) == 3 + sizes = [len(stmt._multi_values[0]) for stmt in conn.statements] + assert sizes == [5, 5, 2] + + seen: list[dict[str, int]] = [] + for stmt in conn.statements: + seen.extend(stmt._multi_values[0]) + assert seen == rows # every row present exactly once, in input order + + +@pytest.mark.unit +def test_no_statement_exceeds_budget_over_columns_rows() -> None: + conn = _FakeConn() + insert_rows(conn, _TABLE, _rows(37), param_budget=10) + for stmt in conn.statements: + assert len(stmt._multi_values[0]) <= 10 // 2 + + +@pytest.mark.unit +def test_each_statement_is_one_multi_row_values_insert_not_executemany() -> None: + # Compiling against the real PG dialect must yield one bind param per + # (row, column) pair -- proof this is a single multi-row VALUES statement, + # not something the driver could page into several round trips. + conn = _FakeConn() + insert_rows(conn, _TABLE, _rows(3), param_budget=10_000) + assert len(conn.statements) == 1 + compiled = conn.statements[0].compile(dialect=pg_dialect()) + assert len(compiled.params) == 3 * 2 + + +@pytest.mark.unit +def test_single_row_still_slices_to_one_statement() -> None: + conn = _FakeConn() + insert_rows(conn, _TABLE, _rows(1), param_budget=10) + assert len(conn.statements) == 1 + assert len(conn.statements[0]._multi_values[0]) == 1 diff --git a/tests/unit/test_chunk_store.py b/tests/unit/test_chunk_store.py index 6be966d..8c089cc 100644 --- a/tests/unit/test_chunk_store.py +++ b/tests/unit/test_chunk_store.py @@ -1,4 +1,4 @@ -"""Unit tests for indexer.chunk_store.write_chunks: statement shape, no DB required. +"""Unit tests for indexer.chunk_store.write_chunks(_batch): statement shape, no DB required. A fake ``Connection`` records the statements/params passed to ``execute`` so the delete-then-insert shape (and the absence of any embedding call) can be asserted @@ -12,7 +12,7 @@ import pytest -from indexer.chunk_store import write_chunks +from indexer.chunk_store import write_chunks, write_chunks_batch class _FakeConn: @@ -37,11 +37,12 @@ def test_deletes_by_file_id_then_inserts_all_rows() -> None: delete_stmt, delete_params = conn.calls[0] assert delete_stmt.table.name == "chunks" - assert delete_params is None + assert delete_params == {"ids": [7]} insert_stmt, values = conn.calls[1] assert insert_stmt.table.name == "chunks" - assert values == [ + assert values is None # rows travel inside stmt.values(...), not as a param list + assert insert_stmt._multi_values[0] == [ { "file_id": 7, "chunk_index": 0, @@ -65,8 +66,8 @@ def test_deletes_by_file_id_then_inserts_all_rows() -> None: def test_ts_column_is_never_written() -> None: conn = _FakeConn() write_chunks(conn, file_id=1, chunks=[(0, "x", 1, 1, [0.0])]) - _insert_stmt, values = conn.calls[1] - assert "ts" not in values[0] + insert_stmt, _values = conn.calls[1] + assert "ts" not in insert_stmt._multi_values[0][0] @pytest.mark.unit @@ -86,3 +87,77 @@ def test_never_touches_an_embedder() -> None: import indexer.chunk_store as chunk_store_module assert "embed" not in vars(chunk_store_module) + + +@pytest.mark.unit +def test_write_chunks_batch_empty_rows_is_a_zero_statement_no_op() -> None: + conn = _FakeConn() + written = write_chunks_batch(conn, rows=[]) + assert written == 0 + assert conn.calls == [] + + +@pytest.mark.unit +def test_write_chunks_batch_deletes_every_id_including_zero_chunk_files() -> None: + conn = _FakeConn() + written = write_chunks_batch( + conn, + rows=[ + (7, [(0, "a", 1, 1, [0.1])]), + (8, []), # zero-chunk file -- still owed the delete + (9, [(0, "b", 2, 2, [0.2]), (1, "c", 3, 3, [0.3])]), + ], + ) + assert written == 3 + assert len(conn.calls) == 2 + + delete_stmt, delete_params = conn.calls[0] + assert delete_stmt.table.name == "chunks" + assert delete_params == {"ids": [7, 8, 9]} + + insert_stmt, _ = conn.calls[1] + file_ids = [row["file_id"] for row in insert_stmt._multi_values[0]] + assert file_ids == [7, 9, 9] # file 8 contributes no insert rows + + +@pytest.mark.unit +def test_write_chunks_batch_duplicate_file_id_keeps_the_last_occurrence( + caplog: pytest.LogCaptureFixture, +) -> None: + """A duplicated file_id must not reach the insert twice. + + Two rows for the same chunks_table.file_id would insert two conflicting + (file_id, chunk_index) tuples -- a real UNIQUE VIOLATION against Postgres, + not just a wasted statement -- so the dedup guard is mandatory, mirroring + indexer.store._flush_file_batch's (path, content_sha) guard. + """ + conn = _FakeConn() + with caplog.at_level("WARNING"): + written = write_chunks_batch( + conn, + rows=[ + (7, [(0, "first", 1, 1, [0.1])]), + (7, [(0, "second", 2, 2, [0.2])]), + ], + ) + assert written == 1 + delete_params = conn.calls[0][1] + assert delete_params == {"ids": [7]} + insert_stmt = conn.calls[1][0] + rows = insert_stmt._multi_values[0] + assert len(rows) == 1 + assert rows[0]["content"] == "second" + assert any("duplicate file_id" in r.message for r in caplog.records) + + +@pytest.mark.unit +def test_write_chunks_and_write_chunks_batch_issue_identical_statements() -> None: + conn_a = _FakeConn() + write_chunks(conn_a, file_id=42, chunks=[(0, "x", 1, 1, [0.0])]) + + conn_b = _FakeConn() + write_chunks_batch(conn_b, rows=[(42, [(0, "x", 1, 1, [0.0])])]) + + assert len(conn_a.calls) == len(conn_b.calls) == 2 + assert conn_a.calls[0][1] == conn_b.calls[0][1] + assert conn_a.calls[1][0]._multi_values[0] == conn_b.calls[1][0]._multi_values[0] diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index f232286..8d1c4a9 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -656,15 +656,18 @@ def fake_embed(texts: list[str]) -> list[list[float]]: # Exercise the closure directly: writing main.py's chunks is a # delete-then-insert against the chunks table, keyed by the given file_id. + # The closure now takes a batch -- a sequence of (file_id, pf) pairs. conn = _FakeChunkConn() pf = ParsedFile(path="main.py", lang="python", size=10, content="def f():\n return 1\n") - idx.chunk_writer(conn, 1, 42, pf) + idx.chunk_writer(conn, 1, [(42, pf)]) assert len(conn.calls) == 2 - delete_stmt, _ = conn.calls[0] + delete_stmt, delete_params = conn.calls[0] assert delete_stmt.table.name == "chunks" + assert delete_params == {"ids": [42]} insert_stmt, values = conn.calls[1] assert insert_stmt.table.name == "chunks" - assert values == [ + assert values is None # rows travel inside stmt.values(...), not as a param list + assert insert_stmt._multi_values[0] == [ { "file_id": 42, "chunk_index": 0, @@ -805,7 +808,7 @@ def test_chunk_writer_covered_guard_skips_an_uncovered_path( chunk_writer = _precompute_chunk_writer([], lambda texts: [[0.0] for _ in texts], 100) conn = _FakeChunkConn() with caplog.at_level(logging.WARNING, logger="indexer.job"): - chunk_writer(conn, 1, 99, pf) + chunk_writer(conn, 1, [(99, pf)]) assert conn.calls == [] assert any( "no precomputed chunks for ghost.py" in r.getMessage() @@ -825,8 +828,8 @@ def test_chunk_writer_covers_every_embedded_path_including_zero_chunk_files() -> empty_pf = ParsedFile(path="empty.py", lang="python", size=0, content="") chunk_writer = _precompute_chunk_writer([empty_pf], lambda texts: [[0.0] for _ in texts], 100) conn = _FakeChunkConn() - chunk_writer(conn, 1, 99, empty_pf) - # write_chunks always issues its DELETE even for zero chunks (see + chunk_writer(conn, 1, [(99, empty_pf)]) + # write_chunks_batch always issues its DELETE even for zero chunks (see # indexer/chunk_store.py) -- so a real (non-warning) statement was issued. assert len(conn.calls) >= 1 diff --git a/tests/unit/test_store_batching.py b/tests/unit/test_store_batching.py new file mode 100644 index 0000000..4992aa3 --- /dev/null +++ b/tests/unit/test_store_batching.py @@ -0,0 +1,522 @@ +"""Unit tests for indexer.store's batched changed/new write path (#105), at the +STATEMENT level. + +Complements ``test_store_delta.py`` (classification) and +``test_store_chunk_writer.py`` (chunk_writer wiring): this module pins the +batch FLUSH itself -- boundary tripping by count and by bytes, the +``excluded.*`` trap, id mapping by ``(path, content_sha)`` (never row order), +the intra-batch dedup guard, and the exact statement inventory of a mixed run. + +Uses the same ``_FakeConn`` idiom as ``test_store_delta.py`` (a hand-rolled +fake ``Connection`` that labels each statement), extended with a +``files_returning``/``membership_returning`` hook so tests can script the +``RETURNING`` row order (or drop a row) independently of insertion order -- +the whole point of tests 7-9 below. +""" + +from __future__ import annotations + +import contextlib +import logging +from collections.abc import Callable +from typing import Any, NamedTuple + +import pytest +from sqlalchemy import Delete, Insert, Update +from sqlalchemy.dialects.postgresql import dialect as pg_dialect + +import indexer.store as store_module +from app.db.models import INDEX_SEMANTICS_VERSION +from indexer.hashing import content_sha +from indexer.languages import ( + ExtractedEdge, + ExtractedSymbol, + FileExtraction, + IndexCounts, + ParsedFile, +) +from indexer.store import index_repo + + +class _IdRow(NamedTuple): + """The batched ``files``/membership-union ``RETURNING id, path, content_sha`` shape.""" + + id: int + path: str + content_sha: str + + +class _FakeResult: + def __init__( + self, + *, + scalar: Any = None, + rowcount: int = 0, + row: Any = None, + rows: list[Any] | None = None, + ) -> None: + self._scalar = scalar + self._row = row + self._rows = rows if rows is not None else [] + self.rowcount = rowcount + + def scalar_one(self) -> Any: + return self._scalar + + def one(self) -> Any: + return self._row + + def all(self) -> list[Any]: + return self._rows + + def __iter__(self) -> Any: + return iter(self._rows) + + +class _FakeConn: + def __init__( + self, + *, + baseline: tuple[str | None, int | None] = (None, None), + carried: set[tuple[str, str]] | None = None, + present: set[tuple[str, str]] | None = None, + provenance: bool = True, + stamp_rowcount: int = 1, + files_returning: Callable[[list[dict[str, Any]]], list[_IdRow]] | None = None, + membership_returning: Callable[[dict[str, Any]], list[_IdRow]] | None = None, + ) -> None: + self._next_file_id = 1 + self._baseline = baseline + self._carried = sorted(carried or set()) + self._present = sorted(present or set()) + self._provenance = provenance + self._stamp_rowcount = stamp_rowcount + self._files_returning = files_returning + self._membership_returning = membership_returning + self.kinds: list[str] = [] + self.files_upsert_stmts: list[Any] = [] + self.symbol_inserts: list[list[dict[str, Any]]] = [] + self.edge_inserts: list[list[dict[str, Any]]] = [] + self.delete_calls: list[tuple[str, list[int]]] = [] + self.membership_params: dict[str, Any] = {} + self.stamp_values: dict[str, Any] = {} + + def begin(self) -> Any: + return contextlib.nullcontext() + + def _text_execute(self, sql: str, params: Any) -> _FakeResult: + if "SELECT path, content_sha FROM files" in sql and "branches @>" in sql: + self.kinds.append("read-carried") + return _FakeResult(rows=[_IdRow(0, p, s) for p, s in self._carried]) + if "SELECT path, content_sha FROM files" in sql: + self.kinds.append("read-present") + return _FakeResult(rows=[_IdRow(0, p, s) for p, s in self._present]) + if "SELECT NOT EXISTS" in sql: + self.kinds.append("provenance-gate") + return _FakeResult(scalar=self._provenance) + if "UPDATE files" in sql and "array_agg(DISTINCT e)" in sql: + self.kinds.append("membership-union") + self.membership_params = dict(params or {}) + if self._membership_returning is not None: + rows = self._membership_returning(params or {}) + else: + rows = [] + for path, sha in zip( + (params or {}).get("paths", []), (params or {}).get("shas", []), strict=True + ): + rows.append(_IdRow(self._next_file_id, path, sha)) + self._next_file_id += 1 + return _FakeResult(rows=rows, rowcount=len(rows)) + if "UPDATE files SET branches = array_remove" in sql: + self.kinds.append("sweep-update") + return _FakeResult(rowcount=0) + if "DELETE FROM files" in sql: + self.kinds.append("sweep-delete") + return _FakeResult(rowcount=0) + raise AssertionError(f"unexpected text() statement: {sql!r}") + + def execute(self, stmt: Any, params: Any = None) -> _FakeResult: + sql = getattr(stmt, "text", None) + if sql is not None: + return self._text_execute(sql, params) + + table = stmt.table.name + if isinstance(stmt, Insert) and table == "repos": + self.kinds.append("repos-insert") + return _FakeResult(scalar=1) + if isinstance(stmt, Insert) and table == "repo_branches": + self.kinds.append("repo-branches-insert") + return _FakeResult(row=self._baseline) + if isinstance(stmt, Insert) and table == "files": + self.kinds.append("files-upsert-batch") + self.files_upsert_stmts.append(stmt) + rows_in = stmt._multi_values[0] + if self._files_returning is not None: + out = self._files_returning(list(rows_in)) + else: + out = [] + for row in rows_in: + file_id = self._next_file_id + self._next_file_id += 1 + out.append(_IdRow(file_id, row["path"], row["content_sha"])) + return _FakeResult(rows=out) + if isinstance(stmt, Delete) and table == "symbols": + self.kinds.append("symbols-delete") + self.delete_calls.append(("symbols", list((params or {}).get("ids", [])))) + return _FakeResult() + if isinstance(stmt, Insert) and table == "symbols": + self.kinds.append("symbols-insert") + self.symbol_inserts.append(list(stmt._multi_values[0])) + return _FakeResult() + if isinstance(stmt, Delete) and table == "reference_edges": + self.kinds.append("edges-delete") + self.delete_calls.append(("reference_edges", list((params or {}).get("ids", [])))) + return _FakeResult() + if isinstance(stmt, Insert) and table == "reference_edges": + self.kinds.append("edges-insert") + self.edge_inserts.append(list(stmt._multi_values[0])) + return _FakeResult() + if isinstance(stmt, Update) and table == "repo_branches": + self.kinds.append("stamp") + self.stamp_values = dict(stmt.compile().params) + return _FakeResult(rowcount=self._stamp_rowcount) + raise AssertionError(f"unexpected statement against {table!r}: {stmt}") + + +def _pf(path: str, content: str) -> ParsedFile: + return ParsedFile(path=path, lang="python", size=len(content.encode()), content=content) + + +def _item( + path: str, content: str, *, symbols: bool = True, edges: bool = False +) -> tuple[ParsedFile, FileExtraction]: + symbol = ExtractedSymbol("f", "function", 1, 2) + return ( + _pf(path, content), + FileExtraction( + symbols=[symbol] if symbols else [], + edges=[ExtractedEdge(kind="call", target="t", line=2, enclosing=symbol)] + if edges + else [], + ), + ) + + +def _key(path: str, content: str) -> tuple[str, str]: + return (path, content_sha(content)) + + +def _index(conn: _FakeConn, items: Any, **kwargs: Any) -> IndexCounts: + return index_repo( + conn, + name="acme/widgets", + branch="main", + is_default=True, + head_sha="sha_new", + items=items, + **kwargs, + ) + + +# --- Test 3: transaction shape pinned across every batching configuration ---- + + +@pytest.mark.unit +@pytest.mark.parametrize("batch_max_files", [1, 2, 7, 500]) +def test_transaction_shape_is_pinned_across_batch_sizes( + monkeypatch: pytest.MonkeyPatch, batch_max_files: int +) -> None: + """T3: repos first, repo_branches second, stamp LAST, sweep immediately + before the stamp -- whatever the batch size.""" + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", batch_max_files) + items = [_item(f"f{i}.py", f"x = {i}\n") for i in range(3)] + conn = _FakeConn() + _index(conn, items) + + assert conn.kinds[0] == "repos-insert" + assert conn.kinds[1] == "repo-branches-insert" + assert conn.kinds[-1] == "stamp" + assert conn.kinds[-3:-1] == ["sweep-update", "sweep-delete"] + + +# --- Test 4: batch boundary by count ------------------------------------------ + + +@pytest.mark.unit +def test_batch_boundary_by_count(monkeypatch: pytest.MonkeyPatch) -> None: + """T4: 1001 changed files at _BATCH_MAX_FILES=500 -> 3 files-upsert-batch + statements of sizes 500/500/1, the third being the post-loop flush.""" + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", 500) + items = [_item(f"f{i}.py", f"x = {i}\n") for i in range(1001)] + conn = _FakeConn() + _index(conn, items) + + sizes = [len(stmt._multi_values[0]) for stmt in conn.files_upsert_stmts] + assert sizes == [500, 500, 1] + + +# --- Test 5: excluded.*, not per-file literals -------------------------------- + + +@pytest.mark.unit +def test_files_upsert_set_clause_uses_excluded_never_a_literal() -> None: + """T5: the SET clause of the batched upsert must bind lang/size/content/commit + via excluded.*, never a Python literal from one file -- the trap that would + make every conflicting row in a batch collapse to the LAST file's values.""" + conn = _FakeConn() + _index(conn, [_item("a.py", "x = 1\n")]) + + stmt = conn.files_upsert_stmts[0] + sql = str(stmt.compile(dialect=pg_dialect())) + set_clause = sql.split("DO UPDATE SET", 1)[1].split("RETURNING", 1)[0] + + assert "excluded.lang" in set_clause + assert "excluded.size" in set_clause + assert "excluded.content" in set_clause + assert "excluded.commit" in set_clause + # No bound literal anywhere in the SET clause -- only the VALUES(...) + # portion (checked separately, not here) may carry bind params. + assert "%(" not in set_clause + + +# --- Test 6: batch boundary by bytes ------------------------------------------ + + +@pytest.mark.unit +def test_batch_boundary_by_bytes_flushes_early_and_never_splits_a_file( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """T6: the byte bound flushes early (post-append check); a single file + individually larger than the bound is neither dropped nor split -- it + flushes with whatever batch it landed in.""" + monkeypatch.setattr(store_module, "_BATCH_MAX_CONTENT_BYTES", 10) + monkeypatch.setattr(store_module, "_BATCH_MAX_FILES", 500) + items = [ + _item("a.py", "aaaaaa"), # 6 bytes + _item("b.py", "bbbbbb"), # 6 bytes -> batch now 12 >= 10, flushes (2 files) + _item("c.py", "c" * 20), # 20 bytes alone, over the bound -- still 1 batch + ] + conn = _FakeConn() + _index(conn, items) + + sizes = [len(stmt._multi_values[0]) for stmt in conn.files_upsert_stmts] + assert sizes == [2, 1] + + +# --- Test 7: chunk_writer once per batch, ids mapped by (path, sha) ---------- + + +@pytest.mark.unit +def test_chunk_writer_called_once_per_batch_with_shuffled_returning_order() -> None: + """T7: ONE chunk_writer call per batch, with ids taken from RETURNING mapped + by (path, content_sha) -- proven by deliberately returning RETURNING rows + in the REVERSE of insertion order.""" + calls: list[list[tuple[int, ParsedFile]]] = [] + + def chunk_writer(conn: Any, repo_id: int, pairs: Any) -> None: + calls.append(list(pairs)) + + def shuffled_returning(rows_in: list[dict[str, Any]]) -> list[_IdRow]: + out = [] + next_id = 100 + for row in reversed(rows_in): + out.append(_IdRow(next_id, row["path"], row["content_sha"])) + next_id += 1 + return out + + conn = _FakeConn(files_returning=shuffled_returning) + items = [_item("a.py", "x = 1\n"), _item("b.py", "y = 2\n"), _item("c.py", "z = 3\n")] + _index(conn, items, chunk_writer=chunk_writer) + + assert len(calls) == 1 + got = {pf.path: file_id for file_id, pf in calls[0]} + # Returned in reverse -> c.py got the FIRST id minted, a.py the LAST. + assert got == {"c.py": 100, "b.py": 101, "a.py": 102} + + +# --- Test 8: a missing RETURNING row raises, no further statement issues ---- + + +@pytest.mark.unit +def test_missing_returning_row_raises_and_issues_no_further_statement() -> None: + """T8: an id missing from RETURNING must RAISE, never warn-and-skip -- a + wrong or silently-dropped file_id here would attach one file's symbols to + another file's row. No delete/insert/stamp follows the failed upsert.""" + + def drop_first(rows_in: list[dict[str, Any]]) -> list[_IdRow]: + out = [] + next_id = 1 + for row in rows_in[1:]: + out.append(_IdRow(next_id, row["path"], row["content_sha"])) + next_id += 1 + return out + + conn = _FakeConn(files_returning=drop_first) + items = [_item("a.py", "x = 1\n"), _item("b.py", "y = 2\n")] + with pytest.raises(RuntimeError, match="a.py"): + _index(conn, items) + + assert "symbols-delete" not in conn.kinds + assert "stamp" not in conn.kinds + + +# --- Test 9: membership row-vanished filter drops just that pair ------------- + + +@pytest.mark.unit +def test_membership_vanished_row_drops_only_that_pair_from_the_chunk_call( + caplog: pytest.LogCaptureFixture, +) -> None: + """T9: the membership seam is already called once with the whole list + (#105 commit 3); a row missing from the union's RETURNING drops only that + one pair (WARNING), the rest of the batch's chunk_writer call is unaffected.""" + calls: list[list[tuple[int, ParsedFile]]] = [] + + def chunk_writer(conn: Any, repo_id: int, pairs: Any) -> None: + calls.append(list(pairs)) + + def only_b(params: dict[str, Any]) -> list[_IdRow]: + # a.py "vanished" between the projection read and the union -- only + # b.py's row comes back from RETURNING. + return [_IdRow(1, "b.py", content_sha("y = 2\n"))] + + items = [_item("a.py", "x = 1\n"), _item("b.py", "y = 2\n")] + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried=set(), + present={_key("a.py", "x = 1\n"), _key("b.py", "y = 2\n")}, + membership_returning=only_b, + ) + with caplog.at_level(logging.WARNING, logger="indexer.store"): + _index(conn, items, chunk_writer=chunk_writer) + + assert len(calls) == 1 + assert [pf.path for _fid, pf in calls[0]] == ["b.py"] + assert any("a.py" in r.getMessage() and "vanished" in r.getMessage() for r in caplog.records) + + +# --- Test 10: intra-batch duplicate (path, content_sha) collapses ----------- + + +@pytest.mark.unit +def test_intra_batch_duplicate_collapses_to_last_wins_with_one_warning( + caplog: pytest.LogCaptureFixture, +) -> None: + """T10: two items sharing (path, content_sha) within one batch collapse to + ONE upsert row (last occurrence wins, matching this module's per-file + last-write-wins) and log exactly one WARNING -- the guard against the + verified ON CONFLICT cardinality-violation.""" + dup_first = _item("a.py", "x = 1\n", symbols=True) + dup_last = _item("a.py", "x = 1\n", symbols=False) # same (path, sha) + conn = _FakeConn() + with caplog.at_level(logging.WARNING, logger="indexer.store"): + counts = _index(conn, [dup_first, dup_last]) + + assert conn.kinds.count("files-upsert-batch") == 1 + assert len(conn.files_upsert_stmts[0]._multi_values[0]) == 1 + duplicate_warnings = [r for r in caplog.records if "duplicate" in r.getMessage()] + assert len(duplicate_warnings) == 1 + # Last occurrence wins: dup_last has no symbols. + assert counts.symbols == 0 + + +# --- Test 11: exact statement inventory over a mixed run -------------------- + + +@pytest.mark.unit +def test_full_statement_inventory_mixed_run_semantic_off( + caplog: pytest.LogCaptureFixture, +) -> None: + """T11 (AC1, unit-tier): 1 unchanged, 1 membership-only, 1 changed + (with edges), 1 new -> the exact inventory. moved.py and new.py flush + TOGETHER in one batch (default _BATCH_MAX_FILES). IndexCounts and the + `delta write set` INFO line are unaffected by batching.""" + unchanged = _item("keep.py", "k = 1\n") + member = _item("shared.py", "s = 1\n") + changed = _item("moved.py", "m = 2\n", edges=True) + added = _item("new.py", "n = 1\n") + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("keep.py", "k = 1\n"), _key("moved.py", "m = 1\n")}, + present={ + _key("keep.py", "k = 1\n"), + _key("moved.py", "m = 1\n"), + _key("shared.py", "s = 1\n"), + }, + ) + with caplog.at_level(logging.INFO, logger="indexer.store"): + counts = _index(conn, [unchanged, member, changed, added]) + + assert conn.kinds == [ + "repos-insert", + "repo-branches-insert", + "read-carried", + "read-present", + "provenance-gate", + "files-upsert-batch", + "symbols-delete", + "symbols-insert", + "edges-delete", + "edges-insert", + "membership-union", + "sweep-update", + "sweep-delete", + "stamp", + ] + assert counts == IndexCounts(files=4, symbols=2, swept=0, edges=1) + assert ( + "acme/widgets@main: delta write set 2/4 files " + "(unchanged=1 membership=1, semantics gate open)" + ) in caplog.text + + +@pytest.mark.unit +def test_full_statement_inventory_mixed_run_semantic_on() -> None: + """T11, semantic-on half: ONE chunk_writer call for the changed/new batch, + a SEPARATE one for the membership class -- never per file.""" + chunk_calls: list[list[str]] = [] + + def chunk_writer(conn: Any, repo_id: int, pairs: Any) -> None: + chunk_calls.append([pf.path for _fid, pf in pairs]) + + unchanged = _item("keep.py", "k = 1\n") + member = _item("shared.py", "s = 1\n") + changed = _item("moved.py", "m = 2\n") + added = _item("new.py", "n = 1\n") + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("keep.py", "k = 1\n"), _key("moved.py", "m = 1\n")}, + present={ + _key("keep.py", "k = 1\n"), + _key("moved.py", "m = 1\n"), + _key("shared.py", "s = 1\n"), + }, + ) + _index(conn, [unchanged, member, changed, added], chunk_writer=chunk_writer) + + assert chunk_calls == [["moved.py", "new.py"], ["shared.py"]] + + +@pytest.mark.unit +def test_empty_membership_within_a_mixed_run_still_skips_the_union() -> None: + """T11, stability half: an empty membership class is still skipped even in + a run that also has a changed/new batch -- keeps the inventories above + stable regardless of what else the run wrote.""" + changed = _item("moved.py", "m = 2\n") + conn = _FakeConn( + baseline=("sha_old", INDEX_SEMANTICS_VERSION), + carried={_key("moved.py", "m = 1\n")}, + present={_key("moved.py", "m = 1\n")}, + ) + _index(conn, [changed]) + assert "membership-union" not in conn.kinds + + +# --- The libpq bind-param invariant ------------------------------------------- + + +@pytest.mark.unit +def test_batch_max_files_times_columns_is_under_the_bind_param_ceiling() -> None: + """libpq's Bind message carries the parameter count in an int16, so no + single statement may bind more than 65535 params.""" + assert store_module._BATCH_MAX_FILES * store_module._FILE_UPSERT_COLUMNS < 65535 diff --git a/tests/unit/test_store_chunk_writer.py b/tests/unit/test_store_chunk_writer.py index a4a788b..670792e 100644 --- a/tests/unit/test_store_chunk_writer.py +++ b/tests/unit/test_store_chunk_writer.py @@ -5,13 +5,14 @@ coverage in tests/integration/test_store.py. This only proves the NEW surface: (a) chunk_writer defaults to None, which is byte-identical to the core path before chunk writing was added, and (b) when given, it is called once per -file, inside the same conn.begin(), with (repo_id, file_id, pf). +BATCH of files, inside the same conn.begin(), with (repo_id, pairs) -- pairs a +sequence of (file_id, pf). """ from __future__ import annotations import contextlib -from typing import Any +from typing import Any, NamedTuple import pytest from sqlalchemy import Delete, Insert, Update @@ -20,10 +21,26 @@ from indexer.store import StaleIndexError, index_repo +class _IdRow(NamedTuple): + """The batched ``files`` upsert's ``RETURNING id, path, content_sha`` shape.""" + + id: int + path: str + content_sha: str + + class _FakeResult: - def __init__(self, *, scalar: Any = None, rowcount: int = 0, row: Any = None) -> None: + def __init__( + self, + *, + scalar: Any = None, + rowcount: int = 0, + row: Any = None, + rows: list[Any] | None = None, + ) -> None: self._scalar = scalar self._row = row + self._rows = rows or [] self.rowcount = rowcount def scalar_one(self) -> Any: @@ -32,6 +49,9 @@ def scalar_one(self) -> Any: def one(self) -> Any: return self._row + def all(self) -> list[Any]: + return self._rows + class _FakeConn: """Just enough of sqlalchemy.Connection for index_repo's fixed statement shape. @@ -40,7 +60,10 @@ class _FakeConn: statement 2 (repo_branches) is the new per-branch CAS baseline; the membership sweep is raw ``text()`` SQL (an UPDATE then a DELETE against ``files``, matched by substring since a ``TextClause`` has no ``.table``); - the final CAS stamp UPDATE targets ``repo_branches``, not ``repos``. + the final CAS stamp UPDATE targets ``repo_branches``, not ``repos``. Since + #105 the ``files`` upsert is a multi-row ``RETURNING id, path, + content_sha``, answered via ``.all()`` and keyed by the batch's own row + values -- not a single ``scalar_one()`` id. """ def __init__(self, *, stamp_rowcount: int = 1) -> None: @@ -66,9 +89,12 @@ def execute(self, stmt: Any, params: Any = None) -> _FakeResult: # (last_indexed_commit, index_semantics_version) -- new-branch shape. return _FakeResult(row=(None, None)) if isinstance(stmt, Insert) and table == "files": - file_id = self._next_file_id - self._next_file_id += 1 - return _FakeResult(scalar=file_id) + rows = [] + for row in stmt._multi_values[0]: + file_id = self._next_file_id + self._next_file_id += 1 + rows.append(_IdRow(file_id, row["path"], row["content_sha"])) + return _FakeResult(rows=rows) if isinstance(stmt, Insert) and table == "symbols": return _FakeResult() if isinstance(stmt, Delete) and table == "symbols": @@ -106,11 +132,12 @@ def test_chunk_writer_defaults_to_none_and_behavior_is_unchanged() -> None: @pytest.mark.unit -def test_chunk_writer_is_called_once_per_file_with_repo_id_and_file_id() -> None: +def test_chunk_writer_is_called_once_per_batch_with_repo_id_and_file_id() -> None: calls: list[tuple[int, int, str]] = [] - def chunk_writer(conn: Any, repo_id: int, file_id: int, pf: ParsedFile) -> None: - calls.append((repo_id, file_id, pf.path)) + def chunk_writer(conn: Any, repo_id: int, pairs: Any) -> None: + for file_id, pf in pairs: + calls.append((repo_id, file_id, pf.path)) items = [ (_pf("a.py", "x = 1\n"), FileExtraction(symbols=[], edges=[])), diff --git a/tests/unit/test_store_delta.py b/tests/unit/test_store_delta.py index 4f90ff0..69c4dc8 100644 --- a/tests/unit/test_store_delta.py +++ b/tests/unit/test_store_delta.py @@ -152,10 +152,13 @@ def execute(self, stmt: Any, params: Any = None) -> _FakeResult: self.kinds.append("repo-branches-insert") return _FakeResult(row=self._baseline) if isinstance(stmt, Insert) and table == "files": - self.kinds.append("file-upsert") - file_id = self._next_file_id - self._next_file_id += 1 - return _FakeResult(scalar=file_id) + self.kinds.append("files-upsert-batch") + rows = [] + for row in stmt._multi_values[0]: + file_id = self._next_file_id + self._next_file_id += 1 + rows.append(_IdRow(file_id, row["path"], row["content_sha"])) + return _FakeResult(rows=rows) if isinstance(stmt, Delete) and table == "symbols": self.kinds.append("symbols-delete") return _FakeResult() @@ -230,7 +233,7 @@ def test_version_mismatch_never_issues_the_projection_reads( assert "read-carried" not in conn.kinds assert "read-present" not in conn.kinds assert "provenance-gate" not in conn.kinds - assert "file-upsert" in conn.kinds + assert "files-upsert-batch" in conn.kinds assert counts == IndexCounts(files=1, symbols=1, swept=0, edges=0) @@ -275,7 +278,7 @@ def test_all_unchanged_run_calls_no_chunk_writer() -> None: _index( conn, items, - chunk_writer=lambda _c, _r, _f, pf: calls.append(pf.path), + chunk_writer=lambda _c, _r, pairs: calls.extend(pf.path for _fid, pf in pairs), ) assert calls == [] @@ -291,7 +294,7 @@ def test_changed_content_at_a_known_path_takes_the_full_path() -> None: ) counts = _index(conn, [_item("a.py", "x = 999\n", edges=True)]) - assert conn.kinds.count("file-upsert") == 1 + assert conn.kinds.count("files-upsert-batch") == 1 assert conn.kinds.count("symbols-delete") == 1 assert conn.kinds.count("edges-delete") == 1 assert counts == IndexCounts(files=1, symbols=1, swept=0, edges=1) @@ -303,24 +306,31 @@ def test_changed_content_at_a_known_path_takes_the_full_path() -> None: @pytest.mark.unit def test_membership_only_issues_one_batched_union_and_no_symbol_work() -> None: """T3 (AC3): a file stored for this repo but not carried by this branch takes - the membership path -- ONE batched UPDATE for the whole class, one - chunk_writer call per file, and no symbols/reference_edges statements.""" - calls: list[tuple[int, str]] = [] + the membership path -- ONE batched UPDATE for the whole class, ONE + chunk_writer call carrying every file's (file_id, pf) pair, and no + symbols/reference_edges statements.""" + chunk_writer_calls: list[list[tuple[int, str]]] = [] items = [_item("a.py", "x = 1\n"), _item("b.py", "y = 2\n")] conn = _FakeConn( baseline=("sha_old", INDEX_SEMANTICS_VERSION), carried=set(), present={_key("a.py", "x = 1\n"), _key("b.py", "y = 2\n")}, ) - counts = _index(conn, items, chunk_writer=lambda _c, _r, fid, pf: calls.append((fid, pf.path))) + counts = _index( + conn, + items, + chunk_writer=lambda _c, _r, pairs: chunk_writer_calls.append( + [(fid, pf.path) for fid, pf in pairs] + ), + ) assert conn.kinds.count("membership-union") == 1 - assert "file-upsert" not in conn.kinds + assert "files-upsert-batch" not in conn.kinds assert "symbols-delete" not in conn.kinds assert "edges-delete" not in conn.kinds - # The file_id each chunk write used came from the UPDATE's RETURNING, not a - # second lookup. - assert calls == [(1, "a.py"), (2, "b.py")] + # ONE chunk_writer call for the whole membership class, and the file_id each + # pair carries came from the UPDATE's RETURNING, not a second lookup. + assert chunk_writer_calls == [[(1, "a.py"), (2, "b.py")]] assert conn.membership_params["paths"] == ["a.py", "b.py"] assert conn.membership_params["branch_arr"] == ["main"] # symbols/edges legitimately fall to zero: no rows were inserted. @@ -343,7 +353,7 @@ def test_membership_only_is_refused_when_a_sibling_branch_is_stale() -> None: assert "provenance-gate" in conn.kinds assert "membership-union" not in conn.kinds - assert conn.kinds.count("file-upsert") == 1 + assert conn.kinds.count("files-upsert-batch") == 1 assert conn.kinds.count("symbols-delete") == 1 assert counts == IndexCounts(files=1, symbols=1, swept=0, edges=0) @@ -351,7 +361,9 @@ def test_membership_only_is_refused_when_a_sibling_branch_is_stale() -> None: @pytest.mark.unit def test_mixed_classification_statement_inventory() -> None: """T5 (AC2): 1 unchanged, 1 membership-only, 1 changed, 1 new -> the exact - inventory, with the batched union issued once, AFTER the per-file loop.""" + inventory. moved.py and new.py are both changed/new, so they flush TOGETHER + in one batch (one files-upsert-batch, one symbols delete+insert, one edges + delete) after the loop; the membership union follows, batched separately.""" unchanged = _item("keep.py", "k = 1\n") member = _item("shared.py", "s = 1\n") changed = _item("moved.py", "m = 2\n") @@ -373,13 +385,8 @@ def test_mixed_classification_statement_inventory() -> None: "read-carried", "read-present", "provenance-gate", - # moved.py -- changed content at a known path - "file-upsert", - "symbols-delete", - "symbols-insert", - "edges-delete", - # new.py -- never seen - "file-upsert", + # moved.py + new.py -- one batch, flushed after the loop + "files-upsert-batch", "symbols-delete", "symbols-insert", "edges-delete", From 5f290c6e19db8614816deeceb3ca74a568ce2ab0 Mon Sep 17 00:00:00 2001 From: Tanner <84605639+IceRhymers@users.noreply.github.com> Date: Sat, 25 Jul 2026 16:49:09 -0700 Subject: [PATCH 7/9] indexer: order-preserving concurrent embedding requests (#107) Refs #107 --- app/AGENTS.md | 2 +- app/config.py | 8 + app/embed.py | 66 +++++- config.yaml | 9 + docs/runbooks/indexing-parallelism.md | 15 ++ docs/runbooks/semantic-enablement.md | 41 ++++ indexer/AGENTS.md | 2 +- indexer/repo_config.py | 10 + scripts/measure_embedding_concurrency.py | 95 ++++++++ tests/unit/test_embed.py | 275 ++++++++++++++++++++++- tests/unit/test_repo_config.py | 12 +- 11 files changed, 521 insertions(+), 14 deletions(-) create mode 100644 scripts/measure_embedding_concurrency.py diff --git a/app/AGENTS.md b/app/AGENTS.md index 6b0867b..07e58b7 100644 --- a/app/AGENTS.md +++ b/app/AGENTS.md @@ -12,7 +12,7 @@ The MCP server Databricks App: a FastMCP streamable-HTTP service exposing the in | `main.py` | FastMCP app: tool registration via `create_app()` factory, per-session `lifespan`, process-scoped engine singleton (`get_engine()`, `threading.Lock` + `atexit` dispose), `anyio` off-loop dispatch under a pool-sized `CapacityLimiter(5)`, `/health` (zero-DB) and `/ready` (`SELECT 1 FROM repos LIMIT 1` grant probe) | | `service.py` | Shared payload builders (`search_code_payload`, `list_repos_payload`, `get_file_payload`, `clamp_limit`): merges grep + symbol legs into the zoekt-parity envelope, base64url pagination cursors (`encode_cursor`/`decode_cursor`, `CursorError` never swallowed), `commit:` prefix resolution against `repo_branches`, permalink-branch selection | | `config.py` | `pydantic-settings` `Settings` with `CODE_SEARCH_` env prefix (timeouts, row limits, semantic tunables); unprefixed `LAKEBASE_ENDPOINT` via `validation_alias`; `SEMANTIC_EMBEDDING_DIM = 1024` single source of truth; `get_settings()` is `lru_cache`d once per process | -| `embed.py` | `EmbedFn` seam (texts → unit-normalized 1024-dim vectors) and `databricks_embedder`: POSTs to the AI Gateway MLflow embeddings route via the SDK's raw API client; lazy SDK import; per-batch count check (`EmbeddingCountMismatchError`) and dim check (`EmbeddingDimMismatchError`) fail loudly instead of misaligning vectors | +| `embed.py` | `EmbedFn` seam (texts → unit-normalized 1024-dim vectors) and `databricks_embedder`: POSTs to the AI Gateway MLflow embeddings route via the SDK's raw API client; lazy SDK import; per-batch count check (`EmbeddingCountMismatchError`) and dim check (`EmbeddingDimMismatchError`) fail loudly instead of misaligning vectors; batches dispatch through an order-preserving `ThreadPoolExecutor.map` (never `as_completed`) at `concurrency` (#107, indexer-only — the query path is one batch and stays serial) | | `app.yaml` | Databricks App runtime config: `ln -sf . app` symlink so `app.` imports resolve at the uploaded working-dir root; shell-form command so `DATABRICKS_APP_PORT` expands; sets only `LAKEBASE_ENDPOINT` | | `requirements.txt` | Deploy-time lockfile exported by `uv export --no-dev --no-hashes --no-emit-project` — regenerate, never hand-edit | | `__init__.py` | Empty package marker | diff --git a/app/config.py b/app/config.py index e6bb63a..1c5322c 100644 --- a/app/config.py +++ b/app/config.py @@ -112,6 +112,14 @@ class Settings(BaseSettings): # which bounds file ingestion, not embedding-chunk granularity. semantic_chunk_max_tokens: int = 512 + # In-flight embedding requests per worker (#107). The indexer clamps to 2 workers when + # semantic is on (indexer/repo_config.py:effective_workers), so total in-flight gateway + # requests are workers x concurrency: 2 x 4 = 8 at this default, 2 x 8 = 16 at the + # config.yaml-enforced ceiling of 8 -- both under the SDK's 20-connection pool + # (pool_block=True, so exceeding it would silently serialize rather than error). Setting + # this to 1 restores today's fully serial embed() and spawns no thread pool. + semantic_embedding_concurrency: int = 4 + @lru_cache(maxsize=1) def get_settings() -> Settings: diff --git a/app/embed.py b/app/embed.py index 2cff250..d65eec3 100644 --- a/app/embed.py +++ b/app/embed.py @@ -18,6 +18,7 @@ from __future__ import annotations from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor from typing import Any from app.config import SEMANTIC_EMBEDDING_DIM, Settings @@ -50,13 +51,27 @@ def _assert_dims(vectors: Sequence[list[float]], dim: int) -> None: def _query_batch( - client: Any, endpoint: str, model: str, batch: list[str], *, max_retries: int + client: Any, + endpoint: str, + model: str, + batch: list[str], + *, + max_retries: int, + ordinal: int = 0, + offset: int = 0, ) -> list[list[float]]: """Query one batch, retrying up to ``max_retries`` times (small, bounded). ``endpoint`` is the gateway route path (e.g. ``/ai-gateway/mlflow/v1/embeddings``), POSTed via the SDK's raw API client; the response is OpenAI-shaped (``{"data": [{"embedding": [...]}, ...]}``). + + ``ordinal``/``offset`` identify this batch's position in the caller's flat text + list (batch index and its starting position, respectively) purely so a count + mismatch can name the offending batch -- serially that batch was recoverable + from context (the one after the last success), but under concurrent dispatch + (:func:`databricks_embedder`) it is not. Both default to 0 so ``_query_batch`` + stays directly callable in isolation. """ last_exc: Exception | None = None for _attempt in range(max_retries + 1): @@ -69,7 +84,8 @@ def _query_batch( # confusing IndexError (or silent corruption) much further downstream. if len(vectors) != len(batch): raise EmbeddingCountMismatchError( - f"embedder returned {len(vectors)} vectors for {len(batch)} texts" + f"embedder returned {len(vectors)} vectors for {len(batch)} texts " + f"(batch {ordinal}, texts[{offset}:{offset + len(batch)}])" ) return vectors except EmbeddingCountMismatchError: @@ -89,6 +105,7 @@ def databricks_embedder( batch_size: int = 64, timeout: float = 20.0, max_retries: int = 2, + concurrency: int = 1, ) -> EmbedFn: """Build an :data:`EmbedFn` backed by the AI Gateway embeddings route ``endpoint``. @@ -105,6 +122,19 @@ def databricks_embedder( client's concern. When omitted, the real ``WorkspaceClient`` is built with a ``Config`` carrying ``http_timeout_seconds=timeout`` (the raw API client has no per-call timeout). + + ``concurrency`` dispatches up to that many batches at once via a + ``ThreadPoolExecutor``. It deliberately does NOT mirror this file's usual + default-mirroring convention (every other parameter here matches its + ``Settings`` twin): the concurrent path must be opt-in at the call site that + knows it is the indexer. ``get_embedder`` supplies the real value from + ``Settings.semantic_embedding_concurrency``; every direct caller and existing + unit test keeps today's serial semantics untouched. Do not "fix" this default + to match -- a higher default would make ``tests/unit/test_embed.py``'s + single-threaded fakes (e.g. ``_FakeApiClient.batches.append``) nondeterministic + under concurrent append. ``concurrency <= 1``, or a text list short enough to + produce only one batch, constructs no pool and spawns no thread -- the query + path (one text, one batch) is a strict no-op. """ if client is None: from databricks.sdk import WorkspaceClient # lazy: see module docstring @@ -113,10 +143,33 @@ def databricks_embedder( client = WorkspaceClient(config=_SdkConfig(http_timeout_seconds=timeout)) def embed(texts: list[str]) -> list[list[float]]: - vectors: list[list[float]] = [] - for i in range(0, len(texts), batch_size): - batch = texts[i : i + batch_size] - vectors.extend(_query_batch(client, endpoint, model, batch, max_retries=max_retries)) + batches = [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)] + workers = max(1, min(concurrency, len(batches))) + + def run_batch(indexed: tuple[int, list[str]]) -> list[list[float]]: + ordinal, batch = indexed + return _query_batch( + client, + endpoint, + model, + batch, + max_retries=max_retries, + ordinal=ordinal, + offset=ordinal * batch_size, + ) + + if workers == 1: + per_batch = [run_batch(item) for item in enumerate(batches)] + else: + # .map() yields in SUBMISSION order regardless of completion order -- + # never as_completed(), which yields in completion order and would + # silently reorder vectors across files. See the module docstring's + # EmbeddingCountMismatchError note and tests/unit/test_embed.py's + # test_embed_module_never_uses_as_completed. + with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="embed") as pool: + per_batch = list(pool.map(run_batch, enumerate(batches))) + + vectors = [v for batch_vectors in per_batch for v in batch_vectors] _assert_dims(vectors, dim) return vectors @@ -140,4 +193,5 @@ def get_embedder(cfg: Settings) -> EmbedFn: dim=cfg.semantic_embedding_dim, batch_size=cfg.semantic_embedding_batch_size, timeout=cfg.semantic_embedding_timeout_s, + concurrency=cfg.semantic_embedding_concurrency, ) diff --git a/config.yaml b/config.yaml index 6094bfa..98e99a7 100644 --- a/config.yaml +++ b/config.yaml @@ -94,6 +94,14 @@ connections: # `enabled: false` makes the job a true semantic no-op (no embedder built, no # chunking, the 2-worker memory clamp not applied) even if the env says enabled — # the fastest way to turn semantic off for the job alone. +# +# `embedding_concurrency` (#107) is in-flight embedding requests PER WORKER, sent +# via a ThreadPoolExecutor that preserves submission order — vectors always come +# back in the order their texts were sent, regardless of which request finishes +# first. Total in-flight gateway requests for the job is workers x concurrency: +# 2 x 4 = 8 at this default, 2 x 8 = 16 at the max of 8, both under the SDK's +# 20-connection pool. Set to 1 to restore fully serial embedding (no thread pool +# spawned at all) if you need to roll back. # semantic: # enabled: true # max_chunks_per_repo: 8000 @@ -101,3 +109,4 @@ connections: # embedding_model: system.ai.gte-large-en # embedding_batch_size: 64 # embedding_timeout_s: 20.0 +# embedding_concurrency: 4 diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index c7e99c7..f9ce380 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/docs/runbooks/indexing-parallelism.md @@ -332,6 +332,21 @@ repo's chunks in memory (~0.5-0.8 GB per worker). The clamp is logged: INFO indexer.job [-]: semantic enabled: clamping index_concurrency 6 -> 2 (memory bound: ...) ``` +### Embedding concurrency (#107) + +`workers x concurrency` is the number that matters, not `concurrency` alone. +Each of the (at most 2, semantic-clamped) workers dispatches up to +`semantic.embedding_concurrency` embedding batches at once +(`app/embed.py:databricks_embedder`, order-preserving `ThreadPoolExecutor.map`): +2 x 4 = 8 in-flight gateway requests at the default, 2 x 8 = 16 at the +config.yaml-enforced ceiling of 8 (the `CODE_SEARCH_SEMANTIC_EMBEDDING_CONCURRENCY` +env var carries no ceiling, mirroring `semantic_embedding_batch_size`'s own +unbounded env surface -- config.yaml is the job's real surface regardless), both +under the SDK's 20-connection pool. +`embedding_concurrency: 1` is the rollback switch — fully serial embedding, no +thread pool spawned. See `docs/runbooks/semantic-enablement.md` §4 for the full +in-flight/memory arithmetic and the 429 posture. + ### The connection pool follows the workers Each worker holds exactly one connection, so the engine is built with diff --git a/docs/runbooks/semantic-enablement.md b/docs/runbooks/semantic-enablement.md index 51ba874..e554afb 100644 --- a/docs/runbooks/semantic-enablement.md +++ b/docs/runbooks/semantic-enablement.md @@ -164,6 +164,46 @@ repo it names. `indexer/job.py` now applies to every index run by default (each worker materialises a whole repo's chunks) — see `docs/runbooks/indexing-parallelism.md`. +**Concurrent embedding requests (#107):** each worker's `embed()` call +(`app/embed.py`) dispatches up to `semantic.embedding_concurrency` batches at once +via a `ThreadPoolExecutor`, using `.map()` — never `as_completed()` — so vectors +always come back in submission order regardless of which request finishes first. +Total in-flight gateway requests for the job is `effective_workers x concurrency`: +2 x 4 = 8 at the default `embedding_concurrency: 4`, 2 x 8 = 16 at the config's +`le=8` ceiling, both under the SDK's 20-connection pool +(`HTTPAdapter(pool_connections=20, pool_maxsize=20, pool_block=True)` — +`pool_block=True` means exceeding the pool **silently serializes** requests rather +than raising, so staying under 20 is the load-bearing bound, not a nice-to-have). +The only new per-in-flight-batch memory cost is transient request/response +buffers (~3.5 MB each: ~2.1 MB parsed vectors + ~1.3 MB raw JSON response + a +small request body) — ~28 MB at the default, ~56 MB at the ceiling, negligible +beside the ~0.5–0.8 GB/worker baseline above. `embedding_concurrency: 1` restores +today's fully serial embedding and spawns no thread pool at all (the rollback +switch). + +429s from the AI Gateway are absorbed entirely by the `databricks-sdk`'s own +`Retry-After`-honouring backoff (`_RetryAfterCustomizer`, defaulting to 1s when +the header is absent) before `_query_batch`'s own bounded retry ever sees them — +`app/embed.py` does not add a third retry layer. That absorption is invisible at +the job's normal INFO log level: the SDK logs each throttle at DEBUG +(`databricks.sdk.retries`). If you suspect throttling during a manual run, set +`logging.getLogger("databricks.sdk.retries").setLevel(logging.DEBUG)` for that +run only — never raise the root logger or the `databricks.sdk` parent logger, +which would also re-enable a request/response body dump (the embedding request +body is repo source code). This is an operator step for a one-off diagnostic +run, never a code change (`tests/unit/test_job_redaction.py` tripwires +`indexer/*.py` against exactly that). + +**Failure-path latency under concurrency:** when one batch raises, the pool's +`__exit__` still waits for every other in-flight request in that worker's pool +to finish before the exception propagates (there is no way to abort an +in-flight HTTP call). A batch-0 failure that returned instantly under serial +dispatch can now wait up to `concurrency - 1` requests' worth of time, each +bounded by `semantic_embedding_timeout_s` (default 20s) plus the SDK's own +retry budget. This is bounded and per-branch, not per-run: a sustained-outage +branch still degrades to a core index without chunks (semantic is additive), +just after a somewhat longer wait than serial dispatch's instant fail-fast. + ## 5. Rollback note `0004`'s `downgrade()` drops the BM25/ANN indexes and the `chunks` table, but **does @@ -192,6 +232,7 @@ semantic: embedding_model: system.ai.gte-large-en embedding_batch_size: 64 embedding_timeout_s: 20.0 + embedding_concurrency: 4 # -> Settings.semantic_embedding_concurrency (#107); 1 = serial ``` Every field is optional; an omitted field falls through to the env value / default, and diff --git a/indexer/AGENTS.md b/indexer/AGENTS.md index 268b661..4477922 100644 --- a/indexer/AGENTS.md +++ b/indexer/AGENTS.md @@ -19,7 +19,7 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` | `job.py` | Entry point and orchestration: `main()`/`run()`, `ThreadPoolExecutor` sized by `effective_workers` (unit of work = one repo, all branches sequential), batched `repo_branches` stamp read, per-branch skip-if-unchanged, per-branch `BranchOutcome` classification (`indexed`/`skipped`/`conflict`/`failed`), semantic precompute (`_precompute_chunk_writer`), `ContextVar`-based `[repo]` log attribution, and per-phase timing instrumentation (one `phase timing repo@branch: total=… resolve=… download=… parse=… embed=… db=… sweep=… other=…` INFO line per INDEXED branch, emitted from inside the worker so `[repo]` resolves; fixed unconditional field set; `_timed_items` charges lazy item production to `parse` so it is not fused into `db`; `resolve=`/`list=` on the per-repo `finished` line carry the repo-scoped costs that belong to no branch's total). Returns 1 if any branch failed; conflicts self-heal and do not fail the run. **Reconciliation checkpoint**: after every worker joins, `_decide_reconciliation` gates on zero failures/conflicts, every repo accounted for, and every repo's `discovery_complete`; on a pass, `_reconcile` reuses the pre-fan-out stamp snapshot to compute each repo's retired branches and the run's full desired repo set, then calls the injected `reconcile_retired_fn`/`reconcile_removed_fn` seams (default: the real `indexer.store` primitives) on one post-fan-out connection — retired branches first, repo purge second. `MAX_PURGE_SHRINK_FRACTION = 0.5` (hardcoded, no config knob) withholds ONLY the purge, as a logged incident signal, if it would remove a strict majority of currently stored repos; retired-branch cleanup on survivors still applies. `ReconcileProgress.committed_any` drives an honest "partially reconciled" vs "left stale" failure message — never the raw exception, only its phase/repo/type name. Never called from worker-thread code (`_index_one*`) — a source-level tripwire test enforces it. | | `languages.py` | Single source of truth shared by `parse.py` and `symbols.py`: `EXT_TO_LANG` (8 extensions -> tree-sitter language names), `SYMBOL_KINDS` (node type -> symbol kind per language), `EDGE_NODE_KINDS` (node type -> reference-edge kind per language, Python-only until #85), `MAX_FILE_BYTES` (1 MB), `SEMANTIC_CHUNK_MAX_CHARS` (2000, ~4 chars/token), and the frozen dataclasses `ParsedFile`, `Chunk`, `ExtractedSymbol`, `ExtractedEdge`, `FileExtraction`, `IndexCounts` (now carries `edges`). | | `parse.py` | `iter_source_files`: **no production caller since #106** — retained deliberately as the extraction ORACLE the streaming path (`ingest.py`) is pinned against by `tests/unit/test_ingest_parity.py`, and undeletable anyway without editing a `SEMANTICS_PATHS` module and forcing an `INDEX_SEMANTICS_VERSION` bump. Walk an extracted tree, yield every text file (unknown extensions kept with `lang=None`, since grep runs over all files); skips `.git/`, symlinks, files > `MAX_FILE_BYTES` (stat before read), NUL-sniffed binaries, and UTF-8 decode failures; strips surviving NULs (Postgres `text` rejects them). `iter_chunks`: deterministic line-aligned chunking, no overlap, no mid-line splits, 1-based inclusive line ranges. | -| `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (0.5 GB peak per worker since #106), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. Deliberately import-light: pydantic + PyYAML + stdlib only. | +| `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (0.5 GB peak per worker since #106), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. `SemanticOverrides.embedding_concurrency` (#107): optional `1..8`, overlays `Settings.semantic_embedding_concurrency` (in-flight embedding requests per worker; job-only, inert on the MCP/webui query path). Deliberately import-light: pydantic + PyYAML + stdlib only. | | `resolve.py` | `resolve_repos`: enumerate org/user selectors, apply `ExcludeRules` to enumerated repos only (explicit `repos` entries always win, unfiltered), dedup case-insensitively keeping first-seen spelling, union each connection's `branches:` globs into per-repo `RepoEntry.branch_globs`, then fail fast: `EmptyConfigError` on zero repos (indexing nothing must not exit 0), `RepoCeilingError` above `MAX_REPOS` (500, overridable via `--max_repos`). Also matches `RepoConfig.semantic_max_chunks_per_repo` case-insensitively onto each resolved repo's `RepoEntry.semantic_max_chunks` (`None` if unmatched); an override key matching no resolved repo logs a WARNING (typo guard). Enumerator params are a test seam only, not provider dispatch. | | `store.py` | `index_repo`: the single atomic unit of work for one (repo, branch) inside `with conn.begin():` — repos upsert, `repo_branches` CAS-baseline read, per-file classification (unchanged / membership-only / changed-new), membership sweep (strip this branch from unseen rows, delete rows with empty `branches`; skipped with a WARNING on an empty seen-set), then CAS stamp (raises `StaleIndexError` on mismatch, rolling everything back). `items` is `Iterable[tuple[ParsedFile, FileExtraction]]`. **Changed/new files are batched (#105)**, not written one at a time: `_flush_file_batch` accumulates up to `_BATCH_MAX_FILES` files or `_BATCH_MAX_CONTENT_BYTES` bytes of `pf.size` (whichever trips first, post-append — a single oversized file is never split), then issues ONE multi-row `files` upsert (`SET` uses `excluded.*` for every per-row column, NEVER a Python literal from one file — the trap that would collapse a batch's conflicting rows onto the last file's values), maps ids back from `RETURNING` by `(path, content_sha)` — never by row order, and RAISES on a missing id rather than warn-and-skip (a wrong id here would attach one file's symbols to another file's row) — then ONE `DELETE ... WHERE file_id = ANY(:ids)` + bulk insert each for `symbols`/`reference_edges` (`reference_edges` unconditional, even when a batch's `FileExtraction.edges` is all-empty, so stale rows never survive a re-index), then ONE `chunk_writer` call for the batch. An intra-batch duplicate `(path, content_sha)` is deduped first (last occurrence wins, one WARNING) — Postgres raises `ON CONFLICT DO UPDATE command cannot affect row a second time` and poisons the transaction otherwise. The sweep call site is wall-clocked into `indexer.timing`'s ambient `PhaseTimer` (`record("sweep", …)`) — never into `IndexCounts` (frozen, compared by value) and never through a new `index_repo` parameter (that signature is an injected seam); `store.py` emits **no** log record for it. Pure DML, no TEMP tables (job role has no TEMP privilege on Lakebase). `reconcile_retired_branches` and `reconcile_removed_repos`: the desired-state storage primitives `job.py`'s reconciliation checkpoint calls — each its own `conn.begin()`, repo-scoped membership subtraction (never delete-by-path) and an exact-match repo purge respectively; neither decides WHAT is retired/desired, only applies a caller-supplied set (see `job.py`'s row). | | `symbols.py` | `extract_file` (#84): one tree-sitter parse, one full-tree walk, emitting both `symbols` and `reference_edges` candidates (`FileExtraction`) in the same pass — nested definitions captured, named nodes only for symbols, kinds from `SYMBOL_KINDS`/`EDGE_NODE_KINDS`, 1-based lines. Stack entries carry the innermost named enclosing definition so edges attribute to it in O(1) with no second walk. `extract_symbols` is a thin wrapper (`extract_file(pf).symbols`) kept for the existing call sites. Python-only edge helpers (`_python_call_target`, `_python_import_edges`) are the seam #85 generalizes to other languages. Parser cache is per-thread (`threading.local`) as insurance against a future GIL-releasing `parse()`. | diff --git a/indexer/repo_config.py b/indexer/repo_config.py index 4d39a9f..d863771 100644 --- a/indexer/repo_config.py +++ b/indexer/repo_config.py @@ -150,6 +150,14 @@ class SemanticOverrides(BaseModel): block only moves the second operand. Setting both is coherent: raise the floor for the whole corpus here, spot-override the outliers in the map. + ``embedding_concurrency`` (#107) bounds in-flight embedding requests per + worker; ``databricks_embedder`` clamps it to the batch count, so a value + larger than the number of batches a repo produces degrades to that count + rather than spawning a thread per batch. The ``le=8`` ceiling here matches + the existing ``index_concurrency`` bounded-field pattern -- see + ``app.config.Settings.semantic_embedding_concurrency`` for the in-flight + arithmetic against the SDK's 20-connection pool. + Two ``Settings`` semantic knobs are deliberately absent, because exposing them would be a lie or a footgun: ``semantic_embedding_dim`` is pinned to ``SEMANTIC_EMBEDDING_DIM`` (the ``chunks.embedding`` column type and the 0004 @@ -180,6 +188,7 @@ class SemanticOverrides(BaseModel): embedding_model: str | None = Field(default=None, min_length=1) embedding_batch_size: int | None = Field(default=None, ge=1) embedding_timeout_s: float | None = Field(default=None, gt=0) + embedding_concurrency: int | None = Field(default=None, ge=1, le=8) @field_validator("embedding_endpoint") @classmethod @@ -232,6 +241,7 @@ def settings_overrides(self) -> dict[str, Any]: "embedding_model": "semantic_embedding_model", "embedding_batch_size": "semantic_embedding_batch_size", "embedding_timeout_s": "semantic_embedding_timeout_s", + "embedding_concurrency": "semantic_embedding_concurrency", } return { settings_field: value diff --git a/scripts/measure_embedding_concurrency.py b/scripts/measure_embedding_concurrency.py new file mode 100644 index 0000000..d72fa0d --- /dev/null +++ b/scripts/measure_embedding_concurrency.py @@ -0,0 +1,95 @@ +"""Measure embedding dispatch scaling under concurrency (#107, AC3 part 1). + +Offline companion to ``app.embed.databricks_embedder``. Drives the SAME dispatch +code path AC3 cares about (batches -> ``ThreadPoolExecutor.map`` -> flatten) +against a fake client whose ``do()`` sleeps a fixed per-batch latency -- no +workspace, no network, not run in CI. It measures that the dispatch actually +parallelizes; it says nothing about the real AI Gateway's throughput or 429 +behavior -- that is the other two halves of AC3's measurement protocol (a +real-repo run's ``embed=`` timing-line comparison, and a single named +``databricks.sdk.retries`` logger), which require a live workspace and are run +by hand, not by this script. See ``docs/runbooks/semantic-enablement.md`` §4. + +Usage: ``uv run python scripts/measure_embedding_concurrency.py +[--latency-s 0.05] [--num-batches 32] [--concurrencies 1,2,4,8]`` +""" + +from __future__ import annotations + +import argparse +import time +from collections.abc import Sequence +from typing import Any + +from app.embed import databricks_embedder + + +class _SleepingApiClient: + """Stands in for ``WorkspaceClient.api_client``: sleeps ``latency_s`` per + batch, then returns one dim-1 vector per text. Order/count correctness is + already pinned by ``tests/unit/test_embed.py``; this script only measures + wall clock.""" + + def __init__(self, latency_s: float) -> None: + self._latency_s = latency_s + + def do(self, method: str, path: str, *, body: dict[str, Any]) -> dict[str, Any]: + time.sleep(self._latency_s) + return {"data": [{"embedding": [0.0]} for _ in body["input"]]} + + +class _SleepingClient: + def __init__(self, latency_s: float) -> None: + self.api_client = _SleepingApiClient(latency_s) + + +def measure(*, num_batches: int, batch_size: int, latency_s: float, concurrency: int) -> float: + """Run one ``embed()`` call over a synthetic corpus of ``num_batches`` + batches and return the wall-clock seconds.""" + texts = [f"text-{i}" for i in range(num_batches * batch_size)] + client = _SleepingClient(latency_s) + embed = databricks_embedder( + "ep", "m", client=client, dim=1, batch_size=batch_size, concurrency=concurrency + ) + start = time.perf_counter() + vectors = embed(texts) + elapsed = time.perf_counter() - start + assert len(vectors) == len(texts) + return elapsed + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--latency-s", type=float, default=0.05, help="fake per-batch latency") + parser.add_argument("--num-batches", type=int, default=32, help="synthetic batch count") + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="texts per batch (kept small so --num-batches controls the batch count directly)", + ) + parser.add_argument( + "--concurrencies", default="1,2,4,8", help="comma-separated concurrency levels" + ) + args = parser.parse_args(argv) + concurrencies = [int(c) for c in args.concurrencies.split(",")] + + print(f"{args.num_batches} batches x {args.latency_s}s fake latency each") + print(f"{'concurrency':>11s} {'wall_clock_s':>12s} {'speedup':>8s}") + baseline: float | None = None + for concurrency in concurrencies: + elapsed = measure( + num_batches=args.num_batches, + batch_size=args.batch_size, + latency_s=args.latency_s, + concurrency=concurrency, + ) + if baseline is None: + baseline = elapsed + speedup = baseline / elapsed if elapsed else float("inf") + print(f"{concurrency:>11d} {elapsed:>12.3f} {speedup:>7.2f}x") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_embed.py b/tests/unit/test_embed.py index f228c58..98cc757 100644 --- a/tests/unit/test_embed.py +++ b/tests/unit/test_embed.py @@ -6,15 +6,21 @@ from __future__ import annotations +import ast +import inspect import sys +import threading +import time from typing import Any import pytest +from app.config import Settings from app.embed import ( EmbeddingCountMismatchError, EmbeddingDimMismatchError, databricks_embedder, + get_embedder, ) @@ -40,13 +46,26 @@ def __init__(self, vectors_fn: Any) -> None: @pytest.mark.unit def test_batching_splits_by_batch_size() -> None: + # concurrency=1 pinned explicitly here: this test's own submission-order + # assertion below would otherwise be the accidental tripwire for the + # default -- passing it explicitly means it no longer is, so + # test_default_concurrency_is_one below covers the default directly. client = _FakeClient(lambda texts: [[0.0, 0.0] for _ in texts]) - embed = databricks_embedder("ep", "m", client=client, dim=2, batch_size=2) + embed = databricks_embedder("ep", "m", client=client, dim=2, batch_size=2, concurrency=1) vectors = embed(["a", "b", "c", "d", "e"]) assert len(vectors) == 5 assert client.api_client.batches == [["a", "b"], ["c", "d"], ["e"]] +@pytest.mark.unit +def test_default_concurrency_is_one() -> None: + """The whole backward-compat guarantee rests on this default: every + non-indexer caller (app/search/semantic.py, every other test in this file) + relies on serial dispatch. Only get_embedder(cfg) is supposed to opt in + to concurrency>1 (see databricks_embedder's docstring).""" + assert inspect.signature(databricks_embedder).parameters["concurrency"].default == 1 + + @pytest.mark.unit def test_single_batch_when_under_batch_size() -> None: client = _FakeClient(lambda texts: [[0.0] for _ in texts]) @@ -121,9 +140,259 @@ def short(texts: list[str]) -> list[list[float]]: @pytest.mark.unit -def test_stub_path_never_imports_databricks_sdk() -> None: +@pytest.mark.parametrize("concurrency", [1, 4]) +def test_stub_path_never_imports_databricks_sdk(concurrency: int) -> None: + # batch_size=1 with two texts -> two batches, so concurrency=4 actually + # reaches the pool branch (concurrent.futures is stdlib; nothing new to + # import). At the module's default batch_size=64 this would be one batch + # and the extension would pin nothing at any concurrency (round-3 finding). sys.modules.pop("databricks.sdk", None) client = _FakeClient(lambda texts: [[0.0, 0.0] for _ in texts]) - embed = databricks_embedder("ep", "m", client=client, dim=2) + embed = databricks_embedder( + "ep", "m", client=client, dim=2, batch_size=1, concurrency=concurrency + ) embed(["hello", "world"]) assert "databricks.sdk" not in sys.modules + + +def _build_order_probe_client() -> tuple[Any, list[str]]: + """Fake client where batch "0" blocks until batch "3" (the last batch) + completes and releases it -- so batch 0 finishes LAST despite being + submitted first. ``completion_order`` records the FAKE's own completion + sequence, for T3 to prove this fixture is not passing vacuously. + """ + release = threading.Event() + completion_order: list[str] = [] + + class _OrderedApiClient: + def do(self, method: str, path: str, *, body: dict[str, Any]) -> dict[str, Any]: + assert method == "POST" + batch = list(body["input"]) + text = batch[0] + if text == "0": + assert release.wait(timeout=5), "the last batch never released batch 0" + completion_order.append(text) + if text == "3": + release.set() + return {"data": [{"embedding": [float(text)]}]} + + class _Client: + def __init__(self) -> None: + self.api_client = _OrderedApiClient() + + return _Client(), completion_order + + +@pytest.mark.unit +def test_concurrent_batches_return_in_submission_order() -> None: + """AC1: vectors come back in submission order even though batch 0 -- the + one that determines index 0 of the result -- is the LAST to complete.""" + client, _ = _build_order_probe_client() + embed = databricks_embedder("ep", "m", client=client, dim=1, batch_size=1, concurrency=4) + vectors = embed(["0", "1", "2", "3"]) + assert vectors == [[0.0], [1.0], [2.0], [3.0]] + + +@pytest.mark.unit +def test_out_of_order_completion_is_recorded_out_of_order() -> None: + """Proves test_concurrent_batches_return_in_submission_order isn't passing + vacuously: the fake's own completion order genuinely differs from submission + order (batch 0 finishes last, not first).""" + client, completion_order = _build_order_probe_client() + embed = databricks_embedder("ep", "m", client=client, dim=1, batch_size=1, concurrency=4) + embed(["0", "1", "2", "3"]) + assert completion_order != ["0", "1", "2", "3"] + assert completion_order[-1] == "0" + + +@pytest.mark.unit +def test_short_batch_under_concurrency_names_the_offending_batch() -> None: + """AC2: a mismatch in a LATE batch under concurrency still names it.""" + + def vectors_fn(batch: list[str]) -> list[list[float]]: + if batch == ["d"]: + return [] # 0 vectors for 1 text: short by one + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder("ep", "m", client=client, dim=1, batch_size=1, concurrency=4) + with pytest.raises(EmbeddingCountMismatchError, match=r"0 vectors for 1 texts \(batch 3"): + embed(["a", "b", "c", "d"]) + + +@pytest.mark.unit +def test_offset_names_the_texts_slice_at_a_non_degenerate_batch_size() -> None: + """T4's ``batch_size=1`` fixture makes ``offset == ordinal`` trivially, so it + cannot catch a broken ``offset = ordinal * batch_size``. This uses a real + batch size (64) so the reported ``texts[192:256]`` only matches if the + multiplication is right.""" + + def vectors_fn(batch: list[str]) -> list[list[float]]: + if len(batch) == 64 and batch[0] == "192": + return [[0.0] for _ in batch[:-1]] # 63 vectors for 64 texts: short by one + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder("ep", "m", client=client, dim=1, batch_size=64, concurrency=4) + texts = [str(i) for i in range(256)] + with pytest.raises( + EmbeddingCountMismatchError, + match=r"63 vectors for 64 texts \(batch 3, texts\[192:256\]\)", + ): + embed(texts) + + +@pytest.mark.unit +def test_first_offending_batch_in_submission_order_wins() -> None: + """A slow low-ordinal failure and a fast high-ordinal failure both occur; + `pool.map` must raise the lowest-ordinal one (deterministic error under + nondeterministic execution), not whichever failed first in wall clock.""" + + def vectors_fn(batch: list[str]) -> list[list[float]]: + text = batch[0] + if text == "1": + time.sleep(0.2) + raise RuntimeError("slow failure at ordinal 1") + if text == "6": + raise RuntimeError("fast failure at ordinal 6") + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder( + "ep", "m", client=client, dim=1, batch_size=1, concurrency=8, max_retries=0 + ) + with pytest.raises(RuntimeError, match="slow failure at ordinal 1"): + embed([str(i) for i in range(8)]) + + +@pytest.mark.unit +def test_failure_cancels_queued_batches() -> None: + """Structural fixture, NOT a scheduler-derived one (see the plan's stall + sweep -- every wall-clock-derived constant fails somewhere). batch_size=1 + with 20 texts is load-bearing: at the module's default batch_size=64, 20 + texts is ONE batch, so this would take the serial path and pass vacuously + (`started == 1 <= 3`) with no pool in either implementation. + + Batch "0" raises immediately; every other batch parks on an Event that is + NEVER set. A parked worker structurally cannot pull the next queue item, so + `started` is fixed by the code path (concurrency slots occupied, plus one + more pulled by the freed failing worker before cancellation fires), not by + scheduling. Measured (per the plan) at exactly 3 with zero variance across + stall levels, while a submit()+serial-.result() refactor starts all 20. + """ + never = threading.Event() + started_lock = threading.Lock() + started = {"n": 0} + + def vectors_fn(batch: list[str]) -> list[list[float]]: + text = batch[0] + with started_lock: + started["n"] += 1 + if text == "0": + raise RuntimeError("endpoint down") + never.wait(timeout=0.5) # parks; never actually set -- bounds runtime only + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + texts = [str(i) for i in range(20)] + embed = databricks_embedder( + "ep", "m", client=client, dim=1, batch_size=1, concurrency=2, max_retries=0 + ) + with pytest.raises(RuntimeError, match="endpoint down"): + embed(texts) + assert started["n"] <= 3 # concurrency (2) + 1, structural, not statistical + + +@pytest.mark.unit +@pytest.mark.parametrize("concurrency", [0, -1, 1]) +def test_concurrency_one_uses_the_calling_thread(concurrency: int) -> None: + """The kill switch is a real no-pool path, and the clamp handles bad + (<=0) values by degrading to serial rather than raising or spawning.""" + caller_thread = threading.current_thread() + seen_threads: list[threading.Thread] = [] + + def vectors_fn(batch: list[str]) -> list[list[float]]: + seen_threads.append(threading.current_thread()) + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder( + "ep", "m", client=client, dim=1, batch_size=1, concurrency=concurrency + ) + embed(["a", "b", "c"]) + assert seen_threads == [caller_thread] * 3 + + +@pytest.mark.unit +def test_single_batch_never_spawns_a_pool() -> None: + """The query path (app/search/semantic.py: one text -> one batch) pays + nothing even at a high configured concurrency.""" + caller_thread = threading.current_thread() + seen_threads: list[threading.Thread] = [] + + def vectors_fn(batch: list[str]) -> list[list[float]]: + seen_threads.append(threading.current_thread()) + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder("ep", "m", client=client, dim=1, concurrency=8) + embed(["only one text"]) + assert seen_threads == [caller_thread] + + +@pytest.mark.unit +def test_dim_mismatch_still_raises_under_concurrency() -> None: + """Aggregate _assert_dims still runs over the flattened result under the + pool path. batch_size=1 with two texts -> two batches, so this actually + exercises the pool branch rather than degrading to serial.""" + client = _FakeClient(lambda texts: [[0.0, 0.0, 0.0] for _ in texts]) # dim 3, expect 2 + embed = databricks_embedder("ep", "m", client=client, dim=2, batch_size=1, concurrency=4) + with pytest.raises(EmbeddingDimMismatchError): + embed(["a", "b"]) + + +@pytest.mark.unit +def test_embed_module_never_uses_as_completed() -> None: + """AST-based, not a substring check: a substring check would fail on this + very module's load-bearing explanatory comment about why as_completed must + never appear here (it yields in completion order, defeating AC1).""" + import app.embed as embed_module + + tree = ast.parse(inspect.getsource(embed_module)) + for node in ast.walk(tree): + if isinstance(node, ast.Name) and node.id == "as_completed": + pytest.fail("app.embed must never reference as_completed -- see module docstring") + if isinstance(node, ast.Attribute) and node.attr == "as_completed": + pytest.fail("app.embed must never reference as_completed -- see module docstring") + if isinstance(node, ast.ImportFrom) and any( + alias.name == "as_completed" for alias in node.names + ): + pytest.fail("app.embed must never import as_completed -- see module docstring") + + +@pytest.mark.unit +def test_get_embedder_threads_concurrency_from_settings(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def fake_databricks_embedder(endpoint: str, model: str, **kwargs: Any) -> Any: + captured.update(kwargs) + return lambda texts: [] + + monkeypatch.setattr("app.embed.databricks_embedder", fake_databricks_embedder) + cfg = Settings(semantic_embedding_endpoint="/ep", semantic_embedding_concurrency=7) + get_embedder(cfg) + assert captured["concurrency"] == 7 + + +@pytest.mark.unit +def test_empty_texts_under_concurrency() -> None: + calls = {"n": 0} + + def vectors_fn(texts: list[str]) -> list[list[float]]: + calls["n"] += 1 + return [[0.0] for _ in texts] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder("ep", "m", client=client, dim=1, concurrency=8) + assert embed([]) == [] + assert calls["n"] == 0 diff --git a/tests/unit/test_repo_config.py b/tests/unit/test_repo_config.py index 51b88b2..e0d0eb4 100644 --- a/tests/unit/test_repo_config.py +++ b/tests/unit/test_repo_config.py @@ -323,7 +323,7 @@ def test_semantic_block_absent_is_a_noop() -> None: @pytest.mark.unit def test_semantic_block_full_maps_every_field_to_its_settings_name() -> None: - """A fully-populated block emits all six Settings-named keys, values intact.""" + """A fully-populated block emits all seven Settings-named keys, values intact.""" raw = ( b"version: 1\nconnections:\n - type: github\n users: [u]\n" b"semantic:\n" @@ -333,6 +333,7 @@ def test_semantic_block_full_maps_every_field_to_its_settings_name() -> None: b" embedding_model: acme.embed-v2\n" b" embedding_batch_size: 32\n" b" embedding_timeout_s: 15.0\n" + b" embedding_concurrency: 4\n" ) cfg = parse_config(raw, source="cfg") assert cfg.semantic.settings_overrides() == { @@ -342,6 +343,7 @@ def test_semantic_block_full_maps_every_field_to_its_settings_name() -> None: "semantic_embedding_model": "acme.embed-v2", "semantic_embedding_batch_size": 32, "semantic_embedding_timeout_s": 15.0, + "semantic_embedding_concurrency": 4, } @@ -371,6 +373,8 @@ def test_semantic_block_partial_emits_only_set_fields() -> None: ("embedding_timeout_s", b"-2.5"), ("embedding_endpoint", b'""'), ("embedding_model", b'""'), + ("embedding_concurrency", b"0"), + ("embedding_concurrency", b"9"), ], ) def test_semantic_block_rejects_out_of_bound_values(field: str, value: bytes) -> None: @@ -471,10 +475,11 @@ def test_settings_overrides_keys_are_real_settings_fields_and_types_survive() -> embedding_model="custom.model", embedding_batch_size=8, embedding_timeout_s=5.5, + embedding_concurrency=6, ) overrides = ov.settings_overrides() - # All six set -> all six emitted, and every key is a real Settings field. - assert len(overrides) == 6 + # All seven set -> all seven emitted, and every key is a real Settings field. + assert len(overrides) == 7 assert set(overrides) <= set(Settings.model_fields) # model_copy(update=) does NOT validate; re-validating the dumped model is what @@ -487,6 +492,7 @@ def test_settings_overrides_keys_are_real_settings_fields_and_types_survive() -> assert revalidated.semantic_embedding_model == "custom.model" assert revalidated.semantic_embedding_batch_size == 8 assert revalidated.semantic_embedding_timeout_s == 5.5 + assert revalidated.semantic_embedding_concurrency == 6 # --- parse failures ------------------------------------------------------- From b6379db13cae8dba249486491f6159f0a2a94e13 Mon Sep 17 00:00:00 2001 From: Tanner <84605639+IceRhymers@users.noreply.github.com> Date: Sun, 26 Jul 2026 09:35:07 -0700 Subject: [PATCH 8/9] indexer: process-pool symbol/edge extraction (#108) Refs #108 --- config.yaml | 16 +- docs/perf/issue-108-measurements.md | 139 ++++++ docs/runbooks/indexing-parallelism.md | 135 +++++- indexer/AGENTS.md | 12 +- indexer/extract_pool.py | 635 ++++++++++++++++++++++++++ indexer/job.py | 71 ++- indexer/repo_config.py | 22 +- scripts/measure_extraction_pool.py | 128 ++++++ tests/unit/test_extract_pool.py | 626 +++++++++++++++++++++++++ tests/unit/test_job.py | 244 +++++++++- tests/unit/test_repo_config.py | 35 ++ 11 files changed, 2020 insertions(+), 43 deletions(-) create mode 100644 docs/perf/issue-108-measurements.md create mode 100644 indexer/extract_pool.py create mode 100644 scripts/measure_extraction_pool.py create mode 100644 tests/unit/test_extract_pool.py diff --git a/config.yaml b/config.yaml index 98e99a7..28eff07 100644 --- a/config.yaml +++ b/config.yaml @@ -8,13 +8,23 @@ version: 1 # disk. That is the only on-disk artifact — the archive is streamed in memory and # never extracted — so budget 0.5 GB per worker: 2 GB at the default 4, 4 GB at # the ceiling of 8. -# Returns at the ceiling are sublinear: symbol extraction does not parallelise -# (measured 0.95x on 4 threads), so you buy far less than 8x for a hard linear -# 4 GB of disk. Raise it knowing that. +# Symbol extraction does not parallelise across THREADS (measured 0.95x on 4 +# threads — the tree walk is GIL-serialized), which is why extraction now runs +# in its own shared process pool instead — see extract_processes below. Raising +# this knob buys disk-bound repo fan-out, not extraction throughput. # With semantic indexing enabled this is clamped to 2 — a MEMORY bound, since # embedding materialises a whole repo's chunks (~0.5-0.8 GB per worker). # index_concurrency: 4 +# How many worker PROCESSES the job uses to extract symbols/edges (issue #108). +# Independent of index_concurrency above: this is a CPU knob (a shared, spawn- +# based process pool decoupled from the per-repo worker threads), not a disk or +# memory one. Default (unset) derives from the runtime's affinity/cgroup-aware +# CPU count, clamped to 8. Setting this to 1 restores fully serial, in-process +# extraction and spawns no process pool at all — the rollback switch if the +# pool ever misbehaves in this runtime. +# extract_processes: 4 + connections: - type: github # orgs / users / repos are UNIONED, then deduplicated by canonical org/repo. diff --git a/docs/perf/issue-108-measurements.md b/docs/perf/issue-108-measurements.md new file mode 100644 index 0000000..463ba65 --- /dev/null +++ b/docs/perf/issue-108-measurements.md @@ -0,0 +1,139 @@ +# Issue #108 — process-pool symbol/edge extraction: measurements + +Two measurement rounds, per the plan's §8.4 requirement. The **planning-time** +probe (§1) established the initial GO decision, against an on-disk-walk baseline +that no longer exists in production. The **shipped** measurement (§2) re-runs the +comparison through the real `indexer.ingest.iter_tar_source_files` tarball source +via `scripts/measure_extraction_pool.py`, and is the number this PR reports as +AC1's evidence — per the plan, it supersedes §1 rather than sitting beside it as +an equal alternative. + +**§2's result requires a stop-and-report, not a silent GO**: see §2.3. + +## 1. Planning-time probe (superseded — on-disk-walk baseline) + +Throwaway `/tmp` probes, not committed to the repo. **Environment:** Linux, 12 +cores, Python 3.12 (`.venv`), `spawn`, 2 MB batches. **Corpus:** this repo +including `.venv/` at planning time — 9605 indexable files / 97.4 MB, 5064 files +/ 65.8 MB qualifying. + +| Path | Wall clock | Speedup | +|---|---|---| +| serial in-process (on-disk walk) | 9.84 s | 1.00x | +| pool, 2 processes | 5.45 s | 1.81x | +| pool, 4 processes | 2.91 s | **3.38x** | +| pool, 8 processes | 2.06 s | 4.78x | + +Output parity `identical=True` at every process count. This number used a +**batch-submit harness** (`executor.map` over a pre-built list), not the shipped +bounded-look-ahead generator, and a source that read from an already-extracted +tree — both superseded by #106's single-pass tarball source. Recorded here for +provenance only; **do not cite as AC1's evidence.** + +## 2. Shipped measurement (real tarball, real `stream()`) + +`scripts/measure_extraction_pool.py --tarball repo.tar.gz --processes 2,4,8 +--repeat 3`, driving the actual `ExtractionPool.stream()` a production branch +calls, over `indexer.ingest.iter_tar_source_files`. + +**Environment:** Linux, 12 cores (`os.sched_getaffinity`; no cgroup-v2 quota on +this box), Python 3.12.13 (`.venv`), `spawn`, 2 MB batches (the default +`_BATCH_BYTES`). + +**Corpus:** a real GitHub-codeload-shaped tarball of this repository's own +working tree at `5f290c6` (the branch point), including `.venv/` as a +site-packages-heavy large-repo proxy (`.venv/bin`'s one absolute-target symlink +excluded — `iter_tar_source_files` correctly rejects it, matching production +behavior for such a member). 4547 indexable files / 54.5 MB, of which **2749 +files / 45.7 MB qualify for parsing** (60% of files, 84% of bytes — a Python- +and-JS-heavy tree). Smaller than the planning-time corpus (the repo has grown +since, and `.venv/bin` is excluded), but the same shape of proxy. + +### 2.1 Result + +``` +$ uv run python scripts/measure_extraction_pool.py --tarball repo.tar.gz --processes 2,4,8 --repeat 3 +loading repo.tar.gz ... +4547 indexable files / 54.5 MB, 2749 qualify for parsing / 45.7 MB (ingest: 0.79s -- serial in the parent, both arms below) + + path extract_s total_s speedup note + serial 6.563 7.355 1.00x + pool x2 3.423 4.215 1.74x identical + pool x4 1.913 2.705 2.72x identical + pool x8 1.422 2.214 3.32x identical + +ingest (serial, shared by both arms): 0.79s of 7.35s serial total (11%) -- the floor AC1's 'combined' speedup above cannot cross, however many processes extract_processes uses. +``` + +Reproduced across three separate invocations (best-of-3 per data point each +time); the pattern is stable, not noise: + +| Run | pool x2 | pool x4 | pool x8 | +|---|---|---|---| +| repeat=1 | 1.81x | 2.75x | 3.22x | +| repeat=3 (a) | 1.73x | 2.75x | 3.22x | +| repeat=3 (b, tabulated above) | 1.74x | **2.72x** | 3.32x | + +Output parity: **`identical=True` at every process count, every run** — the +pooled result list compared element-wise against the serial `extract_file` list. + +### 2.2 Where the two numbers in the table come from + +- **`extract_s`** — the pool's own wall clock (`ExtractionPool.stream()` alone), + comparable to the planning-time probe's number. +- **`total_s`** — `extract_s` plus the ONE shared `ingest` cost (identical in + both arms, since both read the same pre-loaded `ParsedFile` list): this is the + number that maps onto a real branch's `phase timing … parse=…` field, because + `_timed_items` charges `ExtractionPool.stream()`'s pull time — which includes + blocking on a future — entirely to `parse` (proven by + `tests/unit/test_job.py::test_pool_engaged_attributes_stream_production_to_parse`, + T10). **`total_s`'s speedup is AC1's evidence**, not `extract_s`'s. + +`extract_s` alone clears 3x at 4 processes (6.563 / 1.913 = **3.43x**, ~86% +parallel efficiency for the CPU-bound work). `total_s` does not, because ingest +is a **fixed, unavoidably serial** 0.79s added to both the numerator and +denominator (Amdahl's law: at an 11% serial fraction, even the CPU-bound part +scaling perfectly to infinity caps combined speedup at 1/0.11 ≈ 9.1x, and at 4x +*ideal* extraction speedup the combined ceiling is 1/(0.11 + 0.89/4) ≈ 3.03x — +so the measured 2.72x reflects real, good parallel efficiency running into a +structural ceiling, not a bug in this implementation). + +### 2.3 The go/no-go decision: STOP AND REPORT, per §8.4 + +**AC1 ("≥3x on 4+ cores") is NOT met at 4 processes on the shipped measurement: +2.72–2.75x combined, reproduced across three runs.** At 8 processes it clears +the bar (3.22–3.32x), but AC1 asks for the bar to be cleared starting at 4, not +only eventually at 8. + +Per the plan's binding instruction (§8.4, R12), this is the exact condition +under which the executor must **stop and report to the operator rather than +quietly shipping a change whose stated acceptance criterion is a measurement it +missed**, and must not retune the corpus until the number cooperates. It has not +been retuned: this is the corpus described in §2's "Corpus" paragraph, measured +as found. + +This is a real, structural finding, not a defect in the pool's own parallel +efficiency (§2.2's Amdahl arithmetic shows the CPU-bound part scales well). It +is a consequence of #106: the file source is now a single serial +gzip-decompress + decode + filter pass that this pool cannot touch (§4.7.1 of +the plan), and on this corpus that pass is 11% of the serial total — enough to +keep the *combined* number under 3x at 4 processes even though the *extraction* +number alone clears it comfortably (3.43x). The plan names the likely +implication directly: *"it may mean the win now lives in parallelizing +*ingest*, which is a different issue."* + +**This blocks a clean GO claim for AC1 as literally worded.** The design itself +(D1–D8: fork-safety, `BrokenProcessPool` blast-radius containment, the +terminable preflight probe, bounded look-ahead, order preservation) is sound and +independently valuable — a shared process pool doubles-plus extraction +throughput at 4 processes and more than triples it at 8, which is a real win for +any dominant-repo run — but AC1's specific "≥3x on 4+ cores" bar is not met at +the "4" endpoint on this corpus. Reported here rather than adjusted to look +better. + +## 3. Serial-fraction split (for #109) + +The number that tells #109 whether raising `extract_processes` past 4 is worth +anything at all: **ingest is 0.79s of the 7.355s serial total (≈11%) on this +corpus.** #109 should treat ~9x as this pool's asymptotic combined-speedup +ceiling on a similarly-shaped corpus, not the naive "N cores → Nx" expectation. diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index f9ce380..ee82ebd 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/docs/runbooks/indexing-parallelism.md @@ -125,7 +125,11 @@ INFO indexer.job [acme/gadgets]: skipped acme/gadgets@main: already indexed at a Wall-clock for the whole run is bounded below by the single slowest repo, so if one repo dominates, raising `index_concurrency` will not help — that is Amdahl's law asserting itself at the repo level, and the fix is to exclude the repo or -accept the duration. +accept the duration. **Still true for `index_concurrency` specifically** — but +since #108 a single dominant repo is no longer bounded by one thread's `parse` +time: its files parse on every `extract_processes` core via the shared +extraction pool (§3), so raising *that* knob can move the giant's own duration, +even though raising `index_concurrency` cannot. **To decide whether tuning is worth it:** compare the total on the completion line against the sum of the per-repo elapsed times. If the total is already @@ -163,15 +167,29 @@ promise that the field set never changes across releases.) `#104` narrows the **db** and **embed** costs to a branch's actual delta, not its size — but it does NOT touch `parse`: extraction still runs on every file every run (tree-sitter must produce a `FileExtraction` before -`index_repo` can classify it), so an all-unchanged branch on a large repo -still pays its full `download`+`parse` cost. See `indexer.store`'s -`delta write set …` line (below) to tell "this branch is genuinely mostly-new" -from "this branch is mostly-unchanged but still parsing everything" — the -latter is exactly the case `#108` (process-pool extraction) or a future -extraction-skip step would address next. - -Four fields need interpretation before you act on them: - +`index_repo` can classify it, and `index_repo`'s classification happens +*downstream* of extraction — the `FileExtraction` for an unchanged file is +computed and then discarded), so an all-unchanged branch on a large repo still +pays its full `download`+`parse` cost. `#108` (process-pool extraction) makes +that cost cheaper by spreading it across cores — it does NOT skip it: every +file is still parsed every run, including the unchanged ones. See +`indexer.store`'s `delta write set …` line (below) to tell "this branch is +genuinely mostly-new" from "this branch is mostly-unchanged but still parsing +everything on every core" — deferring extraction past classification for the +unchanged fraction is a real, larger follow-up win, deliberately left out of +`#108`'s scope (it would change `index_repo`'s `items` contract). + +Five fields need interpretation before you act on them: + +- **`parse=` can RISE on a single branch after #108, even though the RUN's + total wall clock falls.** The shared extraction pool (§3) is shared across + every `index_concurrency` repo worker, so time a branch spends blocked in + `ExtractionPool.stream()` now includes queueing behind OTHER branches' + batches, not just this branch's own extraction. Reading one branch's `parse=` + in isolation and comparing it to a pre-#108 run will therefore sometimes look + like a regression when it is not — compare **run totals and the dominant + repo** instead (see `docs/perf/issue-108-measurements.md` for the shape of + that comparison). - **`resolve=0.00s` on a default branch is expected, not a bug.** That branch's HEAD SHA came from the repo-level resolve, which happens once per repo outside every branch's total and is reported on the repo's `finished` line as @@ -300,6 +318,7 @@ source edit and redeploy, never a 2am incident-response lever. | `index_concurrency` | 1..8, default **4** | Repos in flight | | `MAX_TARBALL_BYTES` | 500 MB | The compressed download, per worker | | `MAX_EXTRACTED_BYTES` | 2 GB | The streamed uncompressed content, per branch | +| `extract_processes` (#108) | 1..8, default: derived (affinity/cgroup, capped 8) | Symbol/edge extraction worker processes, shared across ALL `index_concurrency` repo workers | Only the first of the two byte caps is a **disk** cap. Since #106 the tarball is streamed once, in memory, and is never extracted, so `MAX_EXTRACTED_BYTES` is a @@ -316,13 +335,17 @@ the compressed tarball is the only artifact on disk, so peak local disk is | **4 (default)** | **2 GB** | | 8 (ceiling) | 4 GB | -**Returns at the ceiling are sublinear; the disk cost is not.** Symbol -extraction was measured at **0.95x on 4 threads** — the tree walk is -GIL-serialized and is ~56% of extraction time, so Amdahl's law caps the speedup -well below 8x. Meanwhile the 4 GB is a hard, linear, unavoidable cost. Raise -`index_concurrency` to 8 only knowing you are buying a fraction of a speedup -with a doubling of disk. (#106 lowered these numbers by 5x but deliberately did -**not** move the default of 4; re-deriving it is #109's job.) +**`index_concurrency` no longer bounds extraction throughput; the disk cost is +still linear.** Symbol extraction was measured at **0.95x on 4 threads** — the +tree walk is GIL-serialized, so raising `index_concurrency` never bought +extraction speedup, only more repos in flight at once. That measurement is +exactly *why* extraction now runs in its own shared **process** pool instead +(`extract_processes`, below) — decoupled from `index_concurrency` entirely. +Raise `index_concurrency` only for repo-level (disk-bound) fan-out; raise +`extract_processes` for CPU-bound extraction throughput. The 4 GB disk figure +above is still a hard, linear, unavoidable cost of `index_concurrency` alone. +(#106 lowered these numbers by 5x but deliberately did **not** move the default +of 4; re-deriving it is #109's job.) **Semantic indexing clamps the pool to 2**, regardless of `index_concurrency`. That clamp is a *memory* bound, not a CPU one: embedding materialises a whole @@ -347,6 +370,84 @@ under the SDK's 20-connection pool. thread pool spawned. See `docs/runbooks/semantic-enablement.md` §4 for the full in-flight/memory arithmetic and the 429 posture. +### Extraction process pool (#108) + +Symbol/edge extraction runs in a **shared, `spawn`-based process pool** — one +pool per run, built once and shared across every `index_concurrency` repo +worker, so a single dominant repo's files parse on every available core instead +of being bound to one thread. This is a **CPU** knob, entirely independent of +`index_concurrency`'s disk bound and the semantic clamp above: it adds no new +corpus writer and does not change per-branch/per-repo sequencing (see §1.1) — +"process pool" reads like a concurrency change to anyone who has internalized +this runbook's §1.1, and it is not one. + +`extract_processes` (unset, the default) derives from the runtime: +affinity/cgroup-aware CPU count, capped at 8 — the same ceiling as every other +parallelism knob here. **Set it to `1` to disable the pool entirely** — fully +serial, in-process extraction, no worker processes spawned at all — mirroring +`embedding_concurrency: 1`'s rollback shape above. This is the 2am escape +hatch if the pool ever misbehaves in this runtime; it is config-only, needs no +redeploy, no migration, and no re-index. + +**How to tell whether the pool engaged**, once per run beside the disk line: + +``` +INFO indexer.extract_pool [-]: symbol extraction: 4 process(es) (spawn); pool preflight ok +``` + +Absent that line (or a `WARNING` in its place, also from the `indexer.extract_pool` +logger), extraction ran in-process — either by config (`extract_processes: 1`) +or because the pool degraded. Three WARNING shapes to recognize: + +- **`extraction pool preflight failed: ...`** — the pool never engaged for this + run at all (a bad sandbox, `/dev/shm` too small, an unguarded `__main__` + under `spawn`). The run proceeds at today's (pre-#108) speed. (A killed probe + worker may also print a stderr line like `resource_tracker: There appear to + be 1 leaked semaphore objects to clean up at shutdown` — harmless noise from + the kill path, not a separate problem.) +- **`... rebuilt the pool (generation N, rebuild M/3)`** — a worker died + (`BrokenProcessPool`, e.g. a native crash in a grammar, an OOM kill). The + branch(es) in flight at that moment failed (up to `index_concurrency` + branches, semantic-clamped to 2 — **not just one**: the pool is shared, so a + break can surface on every repo worker holding a future at that instant). + Each failed branch re-indexes on its next run (it never got a stamp — the + same self-healing property §1 describes). The pool is rebuilt and later + branches run normally. +- **`... rebuild budget (3) exhausted; latching to in-process extraction`** — a + deterministically poisoned file re-broke the pool three times running. The + rest of THIS run finishes in-process (slower, but correct and complete); the + next run gets a fresh pool. + +**Small repos engage fewer workers than `extract_processes` requests, and that +is not a bug.** Batching is by 2 MB of aggregate qualifying-file content, not +file count — a repo with less than `extract_processes x 2 MB` of parseable +content never fills every worker. Such repos are fast regardless; watching one +show fewer active workers than configured does not mean the pool "isn't +working". + +**Ingestion stays serial — the pool cannot parallelize it.** Since #106 the tar +stream is decompressed, decoded, and filtered in ONE pass on the repo-worker +thread (`indexer.ingest.iter_tar_source_files`); the pool only parallelizes +`extract_file` itself. On a real measured corpus this serial pass was ~11% of +the pre-#108 serial total — small, but by Amdahl's law it is a hard ceiling on +the *combined* (ingest + extract) speedup this pool can ever deliver, however +many processes `extract_processes` uses. See +`docs/perf/issue-108-measurements.md` for the full measurement and why it means +AC1's "≥3x on 4+ cores" bar is met at 8 processes but not at 4 on that corpus — +a real, structural finding, not a defect in the pool's own parallel efficiency +(extraction alone scales at ~86% efficiency at 4 processes; the ceiling is the +serial ingest pass, not the pool). + +**Memory**, added to the peak-usage arithmetic #109 inherits: each extraction +worker process peaks at roughly 121 MB RSS (measured, all 7 grammars touched), +so the pool adds `extract_processes x ~121 MB` on top of everything above — +~484 MB at the default-derived 4, ~1.0 GB at the ceiling of 8. The pool's own +bounded look-ahead window (scaled by `2 x extract_processes`: up to 2 MB of +aggregate content per repo worker, plus a file-count backstop for near-zero- +byte files) adds a few tens of MB per in-flight branch on top of that, and +sits alongside (not instead of) the batched-write consumer's own up-to-8-MiB +retention (#105, §2.3) for the same in-flight branch. + ### The connection pool follows the workers Each worker holds exactly one connection, so the engine is built with diff --git a/indexer/AGENTS.md b/indexer/AGENTS.md index 4477922..ab62cb5 100644 --- a/indexer/AGENTS.md +++ b/indexer/AGENTS.md @@ -10,16 +10,17 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` | File | Description | |------|-------------| | `__init__.py` | Empty package marker. | +| `extract_pool.py` | `ExtractionPool` (#108): a shared, `spawn`-based `ProcessPoolExecutor` for symbol/edge extraction — one pool per run, built in `job.run()` and threaded down to every branch, so a dominant repo's files parse on every `extract_processes` core instead of one thread's worth. `stream(files)` batches qualifying files (per `_needs_parse`, mirroring `extract_file`'s own short-circuit) by aggregate content bytes (`_BATCH_BYTES`, 2 MB), bounds look-ahead by two independent budgets scaled by `2 x extract_processes` (pending content bytes, and pending file count — a backstop for near-zero-byte files, which the byte budget alone would not throttle), and yields `(ParsedFile, FileExtraction)` pairs in source order. `extract_processes: 1` (or a failed preflight probe) degrades `stream()` to plain in-process `extract_file` calls — no pool built at all. A `BrokenProcessPool` raises `ExtractionPoolError` (chained), caught by `_index_one_branch`'s existing broad `except`; the supervisor rebuilds once per generation (CAS under a lock, bounded by `MAX_POOL_REBUILDS = 3`) then latches to in-process. The preflight probe (`build_extraction_pool`) is a bare, terminable `multiprocessing.Process` run BEFORE any executor exists — never `executor.submit()` + timeout, which cannot be cancelled and hangs the interpreter at exit. Not watched by the semantics tripwire: it changes *where* extraction runs, never *what* is extracted. | | `branches.py` | Resolves a connection's branch globs against a repo's real branch list via `fnmatchcase`, returning a typed `BranchResolution` (`branches`, `complete`, `dropped`, `cap`). Default branch always included and always first; empty globs means default-branch-only with no GitHub branches API call, always `complete=True`. `SOFT_BRANCH_CAP = 20` truncates (loud warning naming the cap and dropped branches, not a failure; sets `complete=False`, blocking reconciliation for that repo — #61); no override flag — fix the config. | | `bulk.py` | `insert_rows` (#105): param-budgeted multi-row `INSERT` helper shared by `store.py` and `chunk_store.py` — explicit `.values([...])` rather than executemany, so the round-trip count is exactly one statement per slice on every driver. `PARAM_BUDGET` (30,000) and `CHUNK_PARAM_BUDGET` (6,000, payload- not param-bounded — each chunk row carries a 1024-float vector). No-op on an empty row list. | | `chunk_store.py` | `write_chunks_batch` (#105): delete-and-reinsert MANY files' rows in the `chunks` table in one `DELETE ... WHERE file_id = ANY(:ids)` (over every id in the call, including zero-chunk files) plus one param-budgeted bulk insert via `bulk.insert_rows`. `write_chunks` is now a one-element wrapper over it (signature/return unchanged). Takes a live `Connection`; never opens an engine, never calls the embedder — vectors arrive precomputed. No `repo_id` parameter: `chunks` is scoped by `file_id` only. | | `fetch.py` | All GitHub HTTP: paginated org/user repo enumeration (`RepoMeta`), branch listing, `resolve_ref`/`resolve_branch_head` (branch name -> immutable SHA), streamed tarball download capped at `MAX_TARBALL_BYTES` (500 MB), `assert_disk_headroom` (0.5 GB per worker — the compressed tarball is the only artifact written to disk). No extraction: `MAX_EXTRACTED_BYTES` moved to `ingest.py` with #106 and is now a work cap, not a disk cap. `RateLimitError` is deliberately narrow: 429 always; 403 only with `Retry-After` or `X-RateLimit-Remaining: 0` — other 403s are permission failures. | | `ingest.py` | `iter_tar_source_files` (#106): the job's only file source. ONE forward pass over the downloaded tarball (`tarfile.open(mode="r:*")` driven by `tf.next()`; no `getmembers()`, no `extractall()`, nothing written to disk), yielding the same `ParsedFile` stream `parse.iter_source_files` would have yielded from an extracted tree — same filter chain, reusing `parse._looks_binary` by import so the binary-sniff rule cannot fork. `size` is `member.size` (the archive-declared length, which exceeds `len(content)` after NUL stripping); `lang` is the RELATIVE path's lowered suffix. Yields in ARCHIVE order (deterministic per commit; nothing downstream depends on the sequence — see the module docstring's D5 note). `tf.members.clear()` is the FIRST statement of the loop body, since the `TarFile` now outlives `engine.connect()`. Replaces `filter="data"` with explicit, probe-verified rules: raw-`..`-component and absolute member names raise, absolute/escaping link targets raise for hardlinks and symlinks alike, special files and benign hardlinks skip with a WARNING, duplicate names are first-wins, and a zero-member archive raises rather than silently stamping the branch as indexed. `MAX_EXTRACTED_BYTES` (2 GB) is enforced incrementally over pre-filter member sizes. | | `hashing.py` | `content_sha`: canonical SHA-256 hex of content (`None` -> empty string). Single source of truth for `files.content_sha`; must stay byte-identical to the `0003` migration's SQL backfill forever or cross-branch dedup silently breaks (`tests/integration/test_content_sha_parity.py` is the gate). | -| `job.py` | Entry point and orchestration: `main()`/`run()`, `ThreadPoolExecutor` sized by `effective_workers` (unit of work = one repo, all branches sequential), batched `repo_branches` stamp read, per-branch skip-if-unchanged, per-branch `BranchOutcome` classification (`indexed`/`skipped`/`conflict`/`failed`), semantic precompute (`_precompute_chunk_writer`), `ContextVar`-based `[repo]` log attribution, and per-phase timing instrumentation (one `phase timing repo@branch: total=… resolve=… download=… parse=… embed=… db=… sweep=… other=…` INFO line per INDEXED branch, emitted from inside the worker so `[repo]` resolves; fixed unconditional field set; `_timed_items` charges lazy item production to `parse` so it is not fused into `db`; `resolve=`/`list=` on the per-repo `finished` line carry the repo-scoped costs that belong to no branch's total). Returns 1 if any branch failed; conflicts self-heal and do not fail the run. **Reconciliation checkpoint**: after every worker joins, `_decide_reconciliation` gates on zero failures/conflicts, every repo accounted for, and every repo's `discovery_complete`; on a pass, `_reconcile` reuses the pre-fan-out stamp snapshot to compute each repo's retired branches and the run's full desired repo set, then calls the injected `reconcile_retired_fn`/`reconcile_removed_fn` seams (default: the real `indexer.store` primitives) on one post-fan-out connection — retired branches first, repo purge second. `MAX_PURGE_SHRINK_FRACTION = 0.5` (hardcoded, no config knob) withholds ONLY the purge, as a logged incident signal, if it would remove a strict majority of currently stored repos; retired-branch cleanup on survivors still applies. `ReconcileProgress.committed_any` drives an honest "partially reconciled" vs "left stale" failure message — never the raw exception, only its phase/repo/type name. Never called from worker-thread code (`_index_one*`) — a source-level tripwire test enforces it. | +| `job.py` | Entry point and orchestration: `main()`/`run()`, `ThreadPoolExecutor` sized by `effective_workers` (unit of work = one repo, all branches sequential), a shared `ExtractionPool` (#108) built once per run and threaded through `_index_one` → `_index_one_inner` → `_index_one_branch` (`run(..., extraction_pool=None)` injectable, `owns_pool` mirrors `owns_engine`/`owns_http`, shut down in `finally` before `engine.dispose()`), batched `repo_branches` stamp read, per-branch skip-if-unchanged, per-branch `BranchOutcome` classification (`indexed`/`skipped`/`conflict`/`failed`), semantic precompute (`_precompute_chunk_writer`), `ContextVar`-based `[repo]` log attribution, and per-phase timing instrumentation (one `phase timing repo@branch: total=… resolve=… download=… parse=… embed=… db=… sweep=… other=…` INFO line per INDEXED branch, emitted from inside the worker so `[repo]` resolves; fixed unconditional field set; `_timed_items` charges lazy item production to `parse` so it is not fused into `db`; `resolve=`/`list=` on the per-repo `finished` line carry the repo-scoped costs that belong to no branch's total). Returns 1 if any branch failed; conflicts self-heal and do not fail the run. **Reconciliation checkpoint**: after every worker joins, `_decide_reconciliation` gates on zero failures/conflicts, every repo accounted for, and every repo's `discovery_complete`; on a pass, `_reconcile` reuses the pre-fan-out stamp snapshot to compute each repo's retired branches and the run's full desired repo set, then calls the injected `reconcile_retired_fn`/`reconcile_removed_fn` seams (default: the real `indexer.store` primitives) on one post-fan-out connection — retired branches first, repo purge second. `MAX_PURGE_SHRINK_FRACTION = 0.5` (hardcoded, no config knob) withholds ONLY the purge, as a logged incident signal, if it would remove a strict majority of currently stored repos; retired-branch cleanup on survivors still applies. `ReconcileProgress.committed_any` drives an honest "partially reconciled" vs "left stale" failure message — never the raw exception, only its phase/repo/type name. Never called from worker-thread code (`_index_one*`) — a source-level tripwire test enforces it. | | `languages.py` | Single source of truth shared by `parse.py` and `symbols.py`: `EXT_TO_LANG` (8 extensions -> tree-sitter language names), `SYMBOL_KINDS` (node type -> symbol kind per language), `EDGE_NODE_KINDS` (node type -> reference-edge kind per language, Python-only until #85), `MAX_FILE_BYTES` (1 MB), `SEMANTIC_CHUNK_MAX_CHARS` (2000, ~4 chars/token), and the frozen dataclasses `ParsedFile`, `Chunk`, `ExtractedSymbol`, `ExtractedEdge`, `FileExtraction`, `IndexCounts` (now carries `edges`). | | `parse.py` | `iter_source_files`: **no production caller since #106** — retained deliberately as the extraction ORACLE the streaming path (`ingest.py`) is pinned against by `tests/unit/test_ingest_parity.py`, and undeletable anyway without editing a `SEMANTICS_PATHS` module and forcing an `INDEX_SEMANTICS_VERSION` bump. Walk an extracted tree, yield every text file (unknown extensions kept with `lang=None`, since grep runs over all files); skips `.git/`, symlinks, files > `MAX_FILE_BYTES` (stat before read), NUL-sniffed binaries, and UTF-8 decode failures; strips surviving NULs (Postgres `text` rejects them). `iter_chunks`: deterministic line-aligned chunking, no overlap, no mid-line splits, 1-based inclusive line ranges. | -| `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (0.5 GB peak per worker since #106), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. `SemanticOverrides.embedding_concurrency` (#107): optional `1..8`, overlays `Settings.semantic_embedding_concurrency` (in-flight embedding requests per worker; job-only, inert on the MCP/webui query path). Deliberately import-light: pydantic + PyYAML + stdlib only. | +| `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (0.5 GB peak per worker since #106), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. `SemanticOverrides.embedding_concurrency` (#107): optional `1..8`, overlays `Settings.semantic_embedding_concurrency` (in-flight embedding requests per worker; job-only, inert on the MCP/webui query path). `extract_processes` (#108): optional `1..8`, sizes the shared extraction-pool process count (`indexer.extract_pool.derive_process_count`); `None` (default) derives from affinity/cgroup CPU count capped at 8; `1` is the kill switch (no pool built at all). A CPU knob, independent of `index_concurrency`'s disk bound and the semantic clamp. Deliberately import-light: pydantic + PyYAML + stdlib only. | | `resolve.py` | `resolve_repos`: enumerate org/user selectors, apply `ExcludeRules` to enumerated repos only (explicit `repos` entries always win, unfiltered), dedup case-insensitively keeping first-seen spelling, union each connection's `branches:` globs into per-repo `RepoEntry.branch_globs`, then fail fast: `EmptyConfigError` on zero repos (indexing nothing must not exit 0), `RepoCeilingError` above `MAX_REPOS` (500, overridable via `--max_repos`). Also matches `RepoConfig.semantic_max_chunks_per_repo` case-insensitively onto each resolved repo's `RepoEntry.semantic_max_chunks` (`None` if unmatched); an override key matching no resolved repo logs a WARNING (typo guard). Enumerator params are a test seam only, not provider dispatch. | | `store.py` | `index_repo`: the single atomic unit of work for one (repo, branch) inside `with conn.begin():` — repos upsert, `repo_branches` CAS-baseline read, per-file classification (unchanged / membership-only / changed-new), membership sweep (strip this branch from unseen rows, delete rows with empty `branches`; skipped with a WARNING on an empty seen-set), then CAS stamp (raises `StaleIndexError` on mismatch, rolling everything back). `items` is `Iterable[tuple[ParsedFile, FileExtraction]]`. **Changed/new files are batched (#105)**, not written one at a time: `_flush_file_batch` accumulates up to `_BATCH_MAX_FILES` files or `_BATCH_MAX_CONTENT_BYTES` bytes of `pf.size` (whichever trips first, post-append — a single oversized file is never split), then issues ONE multi-row `files` upsert (`SET` uses `excluded.*` for every per-row column, NEVER a Python literal from one file — the trap that would collapse a batch's conflicting rows onto the last file's values), maps ids back from `RETURNING` by `(path, content_sha)` — never by row order, and RAISES on a missing id rather than warn-and-skip (a wrong id here would attach one file's symbols to another file's row) — then ONE `DELETE ... WHERE file_id = ANY(:ids)` + bulk insert each for `symbols`/`reference_edges` (`reference_edges` unconditional, even when a batch's `FileExtraction.edges` is all-empty, so stale rows never survive a re-index), then ONE `chunk_writer` call for the batch. An intra-batch duplicate `(path, content_sha)` is deduped first (last occurrence wins, one WARNING) — Postgres raises `ON CONFLICT DO UPDATE command cannot affect row a second time` and poisons the transaction otherwise. The sweep call site is wall-clocked into `indexer.timing`'s ambient `PhaseTimer` (`record("sweep", …)`) — never into `IndexCounts` (frozen, compared by value) and never through a new `index_repo` parameter (that signature is an injected seam); `store.py` emits **no** log record for it. Pure DML, no TEMP tables (job role has no TEMP privilege on Lakebase). `reconcile_retired_branches` and `reconcile_removed_repos`: the desired-state storage primitives `job.py`'s reconciliation checkpoint calls — each its own `conn.begin()`, repo-scoped membership subtraction (never delete-by-path) and an exact-match repo purge respectively; neither decides WHAT is retired/desired, only applies a caller-supplied set (see `job.py`'s row). | | `symbols.py` | `extract_file` (#84): one tree-sitter parse, one full-tree walk, emitting both `symbols` and `reference_edges` candidates (`FileExtraction`) in the same pass — nested definitions captured, named nodes only for symbols, kinds from `SYMBOL_KINDS`/`EDGE_NODE_KINDS`, 1-based lines. Stack entries carry the innermost named enclosing definition so edges attribute to it in O(1) with no second walk. `extract_symbols` is a thin wrapper (`extract_file(pf).symbols`) kept for the existing call sites. Python-only edge helpers (`_python_call_target`, `_python_import_edges`) are the seam #85 generalizes to other languages. Parser cache is per-thread (`threading.local`) as insurance against a future GIL-releasing `parse()`. | @@ -40,9 +41,10 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - **Test seams are injected callables**, not registries: `run()` accepts `workspace_client`, `http_client`, `engine`, `index_fn`, `embed_fn`, `config_loader`, `reconcile_retired_fn`, `reconcile_removed_fn`; `resolve_repos` accepts enumerators; `store.index_repo`/`reconcile_retired_branches`/`reconcile_removed_repos` and `chunk_store.write_chunks` take a live `Connection`. `resolve_repos` itself is deliberately NOT injectable in `run()` — tests drive it through `httpx.MockTransport`. - **The schema is provider-extensible; the code is GitHub-only.** No `type`-keyed dispatch anywhere — adding GitLab means a new union member plus a new fetch implementation, not a provider abstraction. - **Import discipline**: `repo_config.py` stays import-light (no httpx/SQLAlchemy/SDK at module level); `normalize_repo` is imported from `repo_config`, never from `job` (import cycle); `app.embed`'s databricks-sdk dependency is only imported when semantic is on. +- **Extraction runs in a shared, `spawn`-based process pool (#108)**: workers are pure `ParsedFile -> FileExtraction` — no logging, no DB, no network, and no `indexer.timing.record()` (the ambient `ContextVar` does not cross a process boundary, as `timing.py`'s own docstring states, naming this issue). `stream()`'s source generator (`indexer.ingest.iter_tar_source_files`, a single-pass `TarFile`) is pulled STRICTLY from the calling thread — never add a prefetch/warm-up thread to advance it; a second thread interleaving `tf.next()` calls corrupts a mid-stream `TarFile` silently, not as an exception. A broken pool fails the branches in flight (up to `effective_workers` of them, not just one — the pool is shared), self-heals by rebuilding (generation-CAS'd, bounded), then latches to in-process extraction. Never `multiprocessing.set_start_method(...)` — that mutates global interpreter state a Databricks `python_wheel_task` wrapper may depend on; pass `spawn` as an explicit context object instead. ### Testing Requirements -- `make test` (`pytest -m "unit or observability"`, no external deps): `tests/unit/test_branches.py`, `test_bulk.py`, `test_chunk_store.py`, `test_chunking.py`, `test_edges.py`, `test_fetch.py`, `test_ingest.py`, `test_ingest_parity.py`, `test_job.py`, `test_job_redaction.py`, `test_languages.py`, `test_parse.py`, `test_repo_config.py`, `test_resolve.py`, `test_symbols.py`, `test_store_batching.py`, `test_store_chunk_writer.py`, `test_store_delta.py`, `test_semantics_version_tripwire.py`, `test_timing.py`. +- `make test` (`pytest -m "unit or observability"`, no external deps): `tests/unit/test_branches.py`, `test_bulk.py`, `test_chunk_store.py`, `test_chunking.py`, `test_edges.py`, `test_extract_pool.py`, `test_fetch.py`, `test_ingest.py`, `test_ingest_parity.py`, `test_job.py`, `test_job_redaction.py`, `test_languages.py`, `test_parse.py`, `test_repo_config.py`, `test_resolve.py`, `test_symbols.py`, `test_store_batching.py`, `test_store_chunk_writer.py`, `test_store_delta.py`, `test_semantics_version_tripwire.py`, `test_timing.py`. `test_job.py`'s shared `_config()` helper pins `extract_processes=1` (the pool kill switch) by default, so its ~150 `run()`-driving tests spawn no worker processes; only `test_extract_pool.py` and a handful of dedicated pool-lifecycle tests in `test_job.py` opt into a real pool. - `make test-integration` (needs Postgres): `tests/integration/test_store.py`, `test_store_batching.py` (#105, batch-size parity/excluded-trap/rollback/statement-count — creates no `chunks` table, runs locally), `test_chunk_batching.py` (#105, batched `write_chunks_batch` on bare `vector`, runs locally), `test_store_chunk_writer.py`, `test_content_sha_parity.py` (the hard gate on `hashing.py` vs the migration backfill), `test_reconcile.py` (the `reconcile_retired_branches`/`reconcile_removed_repos` storage primitives directly), and `test_job_reconcile.py` (#59: the real `run()` end to end — real primitives, real engine, only GitHub HTTP faked). - Schema tests must stay fast: changes to `repo_config.py` must not add heavy imports. @@ -51,7 +53,7 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - `pg_insert(...).on_conflict_do_update(...).returning(...)` with a no-op `SET` on conflict — `DO NOTHING ... RETURNING` returns no row on conflict and would break the id bootstrap. - Delete-and-reinsert for child rows with no natural key (symbols, chunks). - `fnmatchcase`, never `fnmatch`: plain globs must behave identically on every platform. -- Guardrail constants with config-level fixes, not override flags: `SOFT_BRANCH_CAP` (20), `MAX_REPOS` (500, the one exception — `--max_repos`), `MAX_FILE_BYTES`, `MAX_TARBALL_BYTES` (`fetch.py`, a disk cap), `MAX_EXTRACTED_BYTES` (`ingest.py` since #106, a WORK cap — nothing is written to disk, so it bounds streamed content, not storage), `MAX_PURGE_SHRINK_FRACTION` (0.5, `job.py` — withholds a corpus-wide repo purge above this shrink fraction; fix is staged config removal, never a knob), `_BATCH_MAX_FILES` (500) / `_BATCH_MAX_CONTENT_BYTES` (8 MiB, `store.py` since #105 — the changed/new write-batch bounds; `_BATCH_MAX_FILES * _FILE_UPSERT_COLUMNS < 65535` is asserted at import time against libpq's Bind param ceiling), `PARAM_BUDGET` (30,000) / `CHUNK_PARAM_BUDGET` (6,000, `bulk.py` since #105 — the generic multi-row-insert slicing bounds). +- Guardrail constants with config-level fixes, not override flags: `SOFT_BRANCH_CAP` (20), `MAX_REPOS` (500, the one exception — `--max_repos`), `MAX_FILE_BYTES`, `MAX_TARBALL_BYTES` (`fetch.py`, a disk cap), `MAX_EXTRACTED_BYTES` (`ingest.py` since #106, a WORK cap — nothing is written to disk, so it bounds streamed content, not storage), `MAX_PURGE_SHRINK_FRACTION` (0.5, `job.py` — withholds a corpus-wide repo purge above this shrink fraction; fix is staged config removal, never a knob), `_BATCH_MAX_FILES` (500) / `_BATCH_MAX_CONTENT_BYTES` (8 MiB, `store.py` since #105 — the changed/new write-batch bounds; `_BATCH_MAX_FILES * _FILE_UPSERT_COLUMNS < 65535` is asserted at import time against libpq's Bind param ceiling), `PARAM_BUDGET` (30,000) / `CHUNK_PARAM_BUDGET` (6,000, `bulk.py` since #105 — the generic multi-row-insert slicing bounds), `_BATCH_BYTES` (2 MB, `extract_pool.py` since #108 — the pool's IPC-amortization batch bound; deliberately `len(pf.content)`, NOT `pf.size` like `store.py`'s batch bound above — one bounds a pickled IPC payload, the other a `files` write payload, and they are allowed to diverge since #106 made `size` exceed `len(content)`) / `MAX_POOL_REBUILDS` (3, `extract_pool.py` — the shared pool's rebuild-then-latch budget). - Errors name the fix: exceptions carry the config key or job parameter to change (`exclude.size_mb`, `index_concurrency`, `--max_repos`). - Broad `except Exception` at classification boundaries only (`run()`'s config/resolve handlers, `_index_one_branch`), each with a comment explaining why broad is load-bearing. - Lazy generators through the open transaction for bounded memory (semantic mode is the exception — it materializes the file list because embedding needs all chunk texts up front). @@ -71,6 +73,6 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` - `pydantic` + `pyyaml` (config schema and parsing) - `tree-sitter` / `tree-sitter-language-pack` (symbol extraction) - `databricks-sdk` (workspace config read + secrets, injected/lazy — never imported at module level here) -- stdlib: `tarfile` (single-pass in-memory streaming in `ingest.py`; no `extractall`, so no `filter="data"` — the rejection rules are explicit there instead), `posixpath` (member-name and link-target safety), `concurrent.futures`, `tempfile`, `shutil`, `fnmatch`, `hashlib`, `argparse`, `contextvars` +- stdlib: `tarfile` (single-pass in-memory streaming in `ingest.py`; no `extractall`, so no `filter="data"` — the rejection rules are explicit there instead), `posixpath` (member-name and link-target safety), `concurrent.futures` (repo-worker thread pool; `ProcessPoolExecutor` for the shared extraction pool since #108), `multiprocessing` (`extract_pool.py`'s `spawn` context and preflight probe, since #108), `tempfile`, `shutil`, `fnmatch`, `hashlib`, `argparse`, `contextvars` diff --git a/indexer/extract_pool.py b/indexer/extract_pool.py new file mode 100644 index 0000000..ee178f6 --- /dev/null +++ b/indexer/extract_pool.py @@ -0,0 +1,635 @@ +"""Shared spawn-based process pool for symbol/edge extraction (issue #108). + +``extract_file`` (``indexer.symbols``) is CPU-bound tree-sitter work that holds +the GIL through ``parse()``, so a repo's own worker THREAD (see ``indexer.job``) +cannot parallelize it. This module decouples extraction from that thread: one +:class:`ExtractionPool` is built once per :func:`indexer.job.run` and shared +across every repo worker, so a single giant repo's files parse on every +available core instead of one thread's worth. + +**Not watched by the semantics tripwire, deliberately.** This module changes +*where* extraction runs, never *what* is extracted -- it imports and calls the +unmodified :func:`indexer.symbols.extract_file` for every file, so output is +identical by construction. It is therefore NOT added to +``tests/unit/test_semantics_version_tripwire.py``'s ``SEMANTICS_PATHS``: that +watch set is for modules that decide extraction *output*, not for a change in +execution topology. See ``indexer/AGENTS.md`` for the ``SEMANTICS_PATHS`` list. + +Four load-bearing design decisions, each with a measured or reproduced failure +mode behind it (see the execution plan for the full writeup): + +* **``spawn``, never ``fork``/``forkserver``.** ``ProcessPoolExecutor`` creates + worker processes lazily on first ``submit()``, so building the pool early does + NOT make ``fork`` safe -- the fork would still happen mid-run, from a process + that by then holds a repo-worker ``ThreadPoolExecutor``, a live SQLAlchemy + engine, an ``httpx.Client``, and (issue #107) a per-embed + ``ThreadPoolExecutor``. Forking a threaded process inherits locks in + indeterminate states -- the classic symptom is a silent stall, not a crash. + ``spawn`` has no inherited-lock or inherited-fd hazard and is passed as an + explicit ``multiprocessing.get_context("spawn")`` object; this module never + calls ``multiprocessing.set_start_method`` (that mutates global interpreter + state a Databricks ``python_wheel_task`` wrapper may depend on). +* **The preflight probe is a bare, terminable ``multiprocessing.Process``, run + BEFORE any ``ProcessPoolExecutor`` exists -- never ``executor.submit()`` + + ``future.result(timeout=...)``.** Reproduced: a timed-out task cannot be + cancelled (``future.cancel()`` returns ``False``, already running), so + ``executor.shutdown(wait=True)`` never returns, and even abandoning the + executor still hangs the interpreter at exit (``concurrent.futures.process`` + registers an atexit hook that joins the executor manager thread). Since + ``resources/job.yml`` sets ``max_concurrent_runs: 1`` with no + ``timeout_seconds``, an executor-based probe that hangs does not degrade to + "a WARNING and a run at today's speed" -- it blocks every queued run + indefinitely until a human intervenes. A bare ``Process`` can be ``kill()``ed + on a timeout, which an executor cannot offer. +* **A broken pool is a per-branch failure, not a run failure, and it + self-heals.** ``ProcessPoolExecutor`` is permanently broken after any worker + dies abnormally: every pending future *and every subsequent submit()* raises + ``BrokenProcessPool``. Because the pool is shared, the true blast radius when + a worker dies is "every repo worker holding a future at that moment" (up to + ``effective_workers`` branches), not one -- state that honestly, not as "one + branch". :class:`ExtractionPool` is a generation-tagged supervisor: it + rebuilds the executor at most once per generation (compare-and-swap on the + generation under a lock, so N concurrent branches hitting the same break + trigger exactly one rebuild), bounded by :data:`MAX_POOL_REBUILDS`. Past that + budget it latches to in-process extraction for the rest of the run and logs a + WARNING -- slower, but correct and complete. +* **No prefetch thread, ever.** Since issue #106 the production file source is + ``indexer.ingest.iter_tar_source_files``, a single forward pass over one open + ``TarFile`` (``tf.members.clear()`` is the first statement of its loop body). + :meth:`ExtractionPool.stream` pulls from that source strictly from the + CALLING thread. A second thread advancing the same generator to "warm" the + next batch would interleave ``tf.next()`` calls on a mid-stream ``TarFile`` + and corrupt output SILENTLY -- not as an exception. Bounded look-ahead + (content-byte and file-count budgets scaled by ``n_processes``, covering + BOTH submitted batches and locally-answered non-qualifying files) is + achieved by pulling more items inside the consumer's own ``next()`` call, + never concurrently. + +No timing, no logging, and no DB/network access happens inside a worker: no +logger is configured in a spawned child (a stray ``logger.info`` there would +silently vanish), and ``indexer.timing``'s ambient ``ContextVar`` does not cross +a process boundary (its own module docstring names this issue by hand). Workers +are pure ``ParsedFile -> FileExtraction``. +""" + +from __future__ import annotations + +import logging +import multiprocessing +import os +import threading +from collections import deque +from collections.abc import Callable, Iterable, Iterator +from concurrent.futures import Future, ProcessPoolExecutor +from concurrent.futures.process import BrokenProcessPool +from pathlib import Path +from typing import TYPE_CHECKING + +from indexer.languages import SYMBOL_KINDS, FileExtraction, ParsedFile +from indexer.symbols import extract_file + +if TYPE_CHECKING: + from indexer.repo_config import RepoConfig + +logger = logging.getLogger("indexer.extract_pool") + +# Aggregate len(pf.content) per submitted batch -- deliberately NOT pf.size +# (indexer/store.py's write-batch bound): since #106 ParsedFile.size is the +# archive-declared length, which exceeds len(content) after NUL stripping. +# store.py bounds a WRITE payload, where the declared size is the honest bound; +# this bounds an IPC payload, which is the string actually pickled. Different +# quantities on purpose -- do not unify them. +_BATCH_BYTES = 2_000_000 + +# Backstop file-count bound, applied at TWO levels: (1) per submitted batch, +# alongside `_BATCH_BYTES` (whichever trips first), and (2) via +# `max_in_flight * _MAX_BATCH_FILES` on the overall look-ahead window -- both +# independent of `_BATCH_BYTES`. Zero- or near-zero-byte files (an empty +# `__init__.py`, a `.gitkeep`) barely move the byte bound, so a corpus +# dominated by them would otherwise defeat it and pickle one enormous batch -- +# mirrors `indexer/store.py`'s own dual byte-and-count bound +# (`_BATCH_MAX_CONTENT_BYTES` / `_BATCH_MAX_FILES`) for the same reason. +_MAX_BATCH_FILES = 2000 + +# Bare-Process preflight probe timeout. Generous on purpose: a healthy runtime +# answers in well under a second (measured ~0.06s); this bound exists only to +# cap the pathological "sandbox silently never reports back" case, and it is +# fully recoverable (kill + latch) either way. +_PROBE_TIMEOUT_S = 60.0 + +# Rebuild-once-per-generation budget for a shared pool that keeps breaking +# (a deterministically poisoned file, a flaky sandbox). Past this, further +# rebuild attempts buy nothing -- latch to in-process and finish the run. +MAX_POOL_REBUILDS = 3 + +_PROBE_SENTINEL = "ok" + + +def _needs_parse(pf: ParsedFile) -> bool: + """Would ``extract_file(pf)`` do any real work, or short-circuit to empty? + + Mirrors ``extract_file``'s own short-circuit (``pf.lang is None`` or the + language has no ``SYMBOL_KINDS`` entry) so the pooled path can answer a + guaranteed-empty file locally instead of paying IPC for it. The two must + stay in exact agreement -- pinned directly by a test that checks every + language in ``EXT_TO_LANG.values()`` plus ``None`` plus a bogus value. + """ + return pf.lang is not None and pf.lang in SYMBOL_KINDS + + +def _extract_batch(batch: list[ParsedFile]) -> list[FileExtraction]: + """Extract every file in ``batch``, in order. Module-level so it is picklable + by qualified name under ``spawn`` -- calls the SAME ``extract_file`` the + in-process path calls, so parity is structural, not re-implemented.""" + return [extract_file(pf) for pf in batch] + + +def _available_cpus() -> int: + """Usable CPU count: affinity-aware, cgroup-v2-quota-aware, floor of 1. + + Prefers ``os.sched_getaffinity(0)`` over ``os.cpu_count()`` (affinity + respects cpuset pinning; ``cpu_count`` reports host cores), then clamps by + the cgroup-v2 CPU quota when ``/sys/fs/cgroup/cpu.max`` is readable and + parses to a smaller number. Every branch here is pure and unit-tested + against synthetic inputs -- this function's REAL return value is never + asserted on in a test (it depends on the test host). + """ + getter = getattr(os, "sched_getaffinity", None) + n = len(getter(0)) if getter is not None else (os.cpu_count() or 1) + quota = _cgroup_cpu_quota() + if quota is not None: + n = min(n, quota) + return max(1, n) + + +def _cgroup_cpu_quota() -> int | None: + """Parse ``/sys/fs/cgroup/cpu.max`` ("$QUOTA $PERIOD" or "max $PERIOD"). + + Returns ``None`` (no clamp) for an unreadable file, an unlimited quota + (``"max"``), or a malformed line -- safe in the correct direction, since + the caller's own ``min(..., 8)`` ceiling still bounds the result. + """ + try: + raw = Path("/sys/fs/cgroup/cpu.max").read_text() + except OSError: + return None + parts = raw.split() + if len(parts) != 2: + return None + quota_str, period_str = parts + if quota_str == "max": + return None + try: + quota = int(quota_str) + period = int(period_str) + except ValueError: + return None + if period <= 0: + return None + return max(1, quota // period) + + +def derive_process_count(config: RepoConfig) -> int: + """Worker-process count for the shared extraction pool. + + ``config.extract_processes`` set explicitly wins outright (``1`` is the + kill switch -- see :func:`build_extraction_pool`). Unset (``None``, the + default) derives from the runtime: affinity/cgroup-aware CPU count, capped + at 8 to match every other parallelism knob in this repo + (``index_concurrency``, ``SemanticOverrides.embedding_concurrency``). + """ + if config.extract_processes is not None: + return config.extract_processes + return min(_available_cpus(), 8) + + +class ExtractionPoolError(RuntimeError): + """A shared extraction pool broke while processing a batch. + + Always chained ``from`` the underlying ``BrokenProcessPool`` and always + fails the branch that raised it -- caught by + ``indexer.job._index_one_branch``'s existing broad ``except Exception``, + exactly like any other extraction failure. The pool itself has already + started (or exhausted its budget for) a rebuild by the time this is + raised; the NEXT branch to call :meth:`ExtractionPool.stream` sees the + rebuilt (or latched-to-in-process) pool, not this one. + """ + + +def _probe_child(queue: "multiprocessing.Queue[str]") -> None: + """Preflight probe worker: exercise the real per-worker import + extraction + path (grammar load included), then report success. Never returns data other + than the bare sentinel -- see :func:`_run_preflight_probe` for why.""" + pf = ParsedFile(path="__preflight__.py", lang="python", size=0, content="def f():\n pass\n") + extract_file(pf) + queue.put(_PROBE_SENTINEL) + + +def _drain(queue: "multiprocessing.Queue[str]") -> str | None: + """Best-effort, non-blocking read of the probe's sentinel. ``None`` on any + failure (empty queue, a queue whose feeder thread never flushed) -- the + caller already knows the child exited 0 by this point; this is corroboration, + not the primary signal.""" + try: + return queue.get_nowait() # type: ignore[no-any-return] + except Exception: + return None + + +def _run_preflight_probe( + *, + target: "Callable[[multiprocessing.Queue[str]], None]" = _probe_child, + timeout: float = _PROBE_TIMEOUT_S, +) -> bool: + """Verify ``spawn``-based multiprocessing actually works here, without ever + constructing a ``ProcessPoolExecutor``. + + A bare, terminable ``multiprocessing.Process`` -- see the module docstring + for why an executor-based probe is actively worse than skipping the probe + entirely. ``join(timeout=...)`` rather than a blocking ``queue.get(timeout= + ...)`` is deliberate: a crashed child is detected the instant it exits, + rather than burning the full timeout waiting on a queue that will never + receive anything. + + ``target``/``timeout`` are test seams only (mirroring this repo's injected- + callable convention -- see ``indexer/AGENTS.md``); production always calls + this with both defaults. Both must be picklable by qualified name under + ``spawn``, so a test override must be a MODULE-LEVEL function, never a + closure or a lambda. + + Every failure mode is caught under a bare ``except Exception``, not a type + list: an unguarded ``__main__`` under ``spawn`` (the serverless hazard this + probe exists to catch -- see the module docstring and ``docs/runbooks``) + makes ``ctx.Process(...).start()`` raise ``RuntimeError`` from + ``multiprocessing.spawn._check_not_importing_main`` in THIS process, not + ``OSError`` -- narrower exception handling here would let that one escape + to the caller instead of latching to in-process extraction. + """ + ctx = multiprocessing.get_context("spawn") + try: + queue: "multiprocessing.Queue[str]" = ctx.Queue() + except Exception: + logger.warning( + "extraction pool preflight failed: could not create an IPC queue " + "(sem_open); falling back to in-process extraction for this run", + exc_info=True, + ) + return False + + process = ctx.Process(target=target, args=(queue,), daemon=True) + try: + process.start() + except Exception: + logger.warning( + "extraction pool preflight failed: could not start a worker process; " + "falling back to in-process extraction for this run", + exc_info=True, + ) + return False + + process.join(timeout=timeout) + if process.is_alive(): + # The disposal an executor cannot offer: kill, then join to reap it. + process.kill() + process.join() + logger.warning( + "extraction pool preflight failed: probe worker did not respond " + "within %.0fs (killed); falling back to in-process extraction for " + "this run", + timeout, + ) + return False + + if process.exitcode != 0: + logger.warning( + "extraction pool preflight failed: probe worker exited with code %s; " + "falling back to in-process extraction for this run", + process.exitcode, + ) + return False + + if _drain(queue) != _PROBE_SENTINEL: + logger.warning( + "extraction pool preflight failed: probe worker exited cleanly but " + "never reported success (missing sentinel); falling back to " + "in-process extraction for this run", + ) + return False + + return True + + +class _FutureHolder: + """One submitted batch's shared future, referenced by every slot in it.""" + + __slots__ = ("future",) + + def __init__(self, future: "Future[list[FileExtraction]]") -> None: + self.future = future + + +class _Slot: + """One file's position in the output stream: either already answered + locally (``local`` set), or awaiting its batch's shared future (``holder`` + + its index within that batch's result list).""" + + __slots__ = ("pf", "local", "holder", "index") + + def __init__(self, pf: ParsedFile, local: FileExtraction | None = None) -> None: + self.pf = pf + self.local = local + self.holder: _FutureHolder | None = None + self.index = 0 + + +class ExtractionPool: + """Supervises a shared, ``spawn``-based extraction pool for one run. + + Owns ``(executor, generation)`` under a lock (D7's supervisor). Built once + per run in ``indexer.job.run`` and shared across every repo worker thread; + ``submit()`` from N concurrent threads into one ``ProcessPoolExecutor`` is + supported (the executor is internally locked). ``n_processes < 2`` means no + executor is ever built -- :meth:`stream` degrades to plain in-process + extraction, exactly today's expression. + """ + + def __init__(self, n_processes: int, *, batch_bytes: int = _BATCH_BYTES) -> None: + self._n_processes = n_processes + self._batch_bytes = batch_bytes + self._lock = threading.Lock() + self._generation = 0 + self._rebuilds = 0 + self._latched = False + self._executor: ProcessPoolExecutor | None = ( + self._new_executor() if n_processes >= 2 else None + ) + + def _new_executor(self) -> ProcessPoolExecutor: + ctx = multiprocessing.get_context("spawn") + return ProcessPoolExecutor(max_workers=self._n_processes, mp_context=ctx) + + def _current(self) -> tuple[ProcessPoolExecutor | None, int]: + with self._lock: + return self._executor, self._generation + + def _handle_broken(self, generation: int) -> None: + """React to a ``BrokenProcessPool`` observed against ``generation``. + + Compare-and-swap on the generation, under the lock, so N concurrent + callers hitting the same break trigger exactly one rebuild (or one + latch): the first caller in advances the generation (or latches); every + other caller sees ``generation != self._generation`` and no-ops. The + caller always raises :class:`ExtractionPoolError` regardless of what + happens here -- this method only updates state for the NEXT call to + :meth:`stream`. + """ + old_executor: ProcessPoolExecutor | None = None + with self._lock: + if self._latched or generation != self._generation: + pass + else: + old_executor = self._executor + if self._rebuilds >= MAX_POOL_REBUILDS: + self._latched = True + self._executor = None + logger.warning( + "extraction pool: rebuild budget (%d) exhausted after " + "repeated BrokenProcessPool; latching to in-process " + "extraction for the rest of this run (rollback: set " + "extract_processes: 1 in config.yaml)", + MAX_POOL_REBUILDS, + ) + else: + self._rebuilds += 1 + self._generation += 1 + self._executor = self._new_executor() + logger.warning( + "extraction pool: a worker died (BrokenProcessPool); " + "rebuilt the pool (generation %d, rebuild %d/%d)", + self._generation, + self._rebuilds, + MAX_POOL_REBUILDS, + ) + # Disposed OUTSIDE the lock, and only with plain shutdown(wait=False): + # a BROKEN executor's manager thread has already terminated, so this + # cannot hang (unlike the healthy-executor teardown in shutdown()). + if old_executor is not None: + old_executor.shutdown(wait=False) + + def stream(self, files: Iterable[ParsedFile]) -> Iterator[tuple[ParsedFile, FileExtraction]]: + """Yield ``(pf, FileExtraction)`` pairs for every file in ``files``, in + source order. + + Pull-driven, strictly from the calling thread: ``files`` (in production, + ``indexer.ingest.iter_tar_source_files``'s single-pass generator) is + never advanced from any other thread. Qualifying files (see + :func:`_needs_parse`) are accumulated into batches bounded by + ``batch_bytes`` of aggregate ``len(pf.content)`` and submitted to the + pool; non-qualifying files are answered locally with + ``FileExtraction([], [])`` and interleaved back at their original + position. Look-ahead is bounded by TWO independent budgets scaled by + ``2 * n_processes`` -- aggregate pending content bytes, and pending + file count (a backstop for a corpus of near-zero-byte files, which the + byte budget alone would not throttle) -- covering both submitted + batches and locally-answered files; pulling more is throttled by + consuming (blocking on) the oldest batch first. + + Raises :class:`ExtractionPoolError` (chained from the underlying + ``BrokenProcessPool``) the first time a submitted batch's worker turns + out to have died -- the caller (one repo branch) fails; every future + call to :meth:`stream` sees whatever :meth:`_handle_broken` decided. + """ + executor, generation = self._current() + if executor is None: + for pf in files: + yield pf, extract_file(pf) + return + yield from self._stream_pooled(files, executor, generation) + + def _stream_pooled( + self, + files: Iterable[ParsedFile], + executor: ProcessPoolExecutor, + generation: int, + ) -> Iterator[tuple[ParsedFile, FileExtraction]]: + # Bounds the total content bytes pulled-but-not-yet-yielded, across + # BOTH local (non-qualifying) and batched slots -- not "in-flight + # batch count". A corpus dominated by non-qualifying files (markdown, + # data, unknown extensions) never submits a batch at all, so a count + # of in-flight batches alone would never throttle it and `pending` + # would grow to the full corpus. Same total budget as a pure + # batch-count scheme (`max_in_flight` batches of up to `_batch_bytes` + # each), just measured directly in bytes so it also covers the + # local-only case. + # + # `max_pending_slots` is a second, independent bound on FILE COUNT: + # zero- or near-zero-byte files (an empty `__init__.py`, a `.gitkeep`) + # contribute ~nothing to `pending_bytes`, so the byte bound alone never + # throttles a corpus dominated by them. Mirrors `indexer/store.py`'s + # own dual byte-and-count bound (`_BATCH_MAX_CONTENT_BYTES` / + # `_BATCH_MAX_FILES`) for the same reason: whichever limit is tighter + # for a given corpus shape is the one that actually binds. + max_in_flight = max(1, 2 * self._n_processes) + max_pending_bytes = max_in_flight * self._batch_bytes + max_pending_slots = max_in_flight * _MAX_BATCH_FILES + source = iter(files) + pending: deque[_Slot] = deque() + in_flight: deque[_FutureHolder] = deque() + pending_bytes = 0 + batch_files: list[ParsedFile] = [] + batch_slots: list[_Slot] = [] + batch_bytes = 0 + exhausted = False + + def flush_batch() -> None: + nonlocal batch_files, batch_slots, batch_bytes + if not batch_files: + return + try: + future = executor.submit(_extract_batch, batch_files) + except BrokenProcessPool as exc: + self._handle_broken(generation) + raise ExtractionPoolError( + f"shared extraction pool (generation {generation}) broke while " + "submitting a batch -- a worker likely crashed or was OOM-killed. " + "This branch failed and will re-index on the next run. Rollback: " + "set extract_processes: 1 in config.yaml to disable pooled " + "extraction." + ) from exc + holder = _FutureHolder(future) + for i, slot in enumerate(batch_slots): + slot.holder = holder + slot.index = i + in_flight.append(holder) + batch_files = [] + batch_slots = [] + batch_bytes = 0 + + def pull_one() -> bool: + """Pull exactly one item from ``source``. ``False`` at EOF (which + also flushes any still-accumulating batch, so nothing is stranded + unsubmitted when the source runs dry).""" + nonlocal batch_bytes, pending_bytes, exhausted + try: + pf = next(source) + except StopIteration: + exhausted = True + flush_batch() + return False + pending_bytes += len(pf.content) + if not _needs_parse(pf): + pending.append(_Slot(pf, local=FileExtraction(symbols=[], edges=[]))) + return True + slot = _Slot(pf) + pending.append(slot) + batch_files.append(pf) + batch_slots.append(slot) + batch_bytes += len(pf.content) + if batch_bytes >= self._batch_bytes or len(batch_files) >= _MAX_BATCH_FILES: + flush_batch() + return True + + try: + while True: + while ( + not exhausted + and pending_bytes < max_pending_bytes + and len(pending) < max_pending_slots + ): + pull_one() + # Force out a still-accumulating partial batch ONLY when it is + # about to be popped next (its first slot is the front of + # `pending`): a long run of non-qualifying (local) files pulled + # AFTER this batch's own files can exhaust the budget while the + # batch itself is still under its own _batch_bytes threshold, + # and since the batch's slots were pulled BEFORE those local + # files, they sit AHEAD of them in `pending` -- reaching the + # front, unflushed, before the local files do. Guarding on + # "is it actually the front" (rather than flushing + # unconditionally every outer-loop iteration) keeps a batch + # accumulating normally toward its full _batch_bytes whenever + # something else is ahead of it in the queue, so this does not + # fragment steady-state batching into one-file submissions. + if pending and pending[0].local is None and pending[0].holder is None: + flush_batch() + if not pending: + return + + slot = pending.popleft() + pending_bytes -= len(slot.pf.content) + if slot.local is not None: + yield slot.pf, slot.local + continue + + holder = slot.holder + # Invariant: every non-local slot has a holder by this point, + # because the unconditional flush_batch() call directly above + # empties `batch_files`/`batch_slots` before we ever pop from + # `pending` -- there is never an unflushed batch slot inside + # `pending` at the moment we reach this line. + assert holder is not None, "unflushed batch slot reached the front of pending" + try: + results = holder.future.result() + except BrokenProcessPool as exc: + self._handle_broken(generation) + raise ExtractionPoolError( + f"shared extraction pool (generation {generation}) broke while " + "processing a batch -- a worker likely crashed or was OOM-killed. " + "This branch failed and will re-index on the next run. Rollback: " + "set extract_processes: 1 in config.yaml to disable pooled " + "extraction." + ) from exc + finally: + if in_flight and in_flight[0] is holder: + in_flight.popleft() + yield slot.pf, results[slot.index] + finally: + # Cleanup on abandonment: `files` is consumed inside index_repo's + # open transaction, so any exception in that loop (a DB error, + # StaleIndexError, a chunk_writer failure) abandons this generator + # with however many batches the look-ahead budgets above still + # left submitted and unconsumed. Cancel every + # one we haven't consumed -- a future already running cannot be + # cancelled (returns False) but finishes in milliseconds and its + # discarded result costs nothing. + for holder in in_flight: + holder.future.cancel() + + def shutdown(self) -> None: + """Tear down the executor, if one exists. ``wait=True``, and safe to + call unconditionally: by the time ``indexer.job.run`` reaches its + ``finally``, the repo-worker ``ThreadPoolExecutor``'s own ``with`` block + has already exited (every branch joined), so no branch holds a future + against this pool -- the one condition that makes ``wait=True`` safe + (see the module's execution-plan writeup). Never + ``cancel_futures=True``, matching this repo's existing executor-teardown + discipline (``indexer/job.py``). + """ + with self._lock: + executor = self._executor + self._executor = None + self._latched = True + if executor is not None: + executor.shutdown(wait=True) + + +def build_extraction_pool(config: RepoConfig) -> ExtractionPool: + """Derive this run's process count, preflight it, and build the pool. + + ``extract_processes: 1`` (explicit or, on a single-core host, derived) is + the kill switch: no preflight probe runs and no executor is ever built -- + :meth:`ExtractionPool.stream` degrades straight to in-process extraction. + For ``>= 2``, the preflight probe (see :func:`_run_preflight_probe`) MUST + pass before any ``ProcessPoolExecutor`` is constructed; a failing probe logs + a WARNING and returns an already-latched, in-process-only pool instead of + raising -- a degraded run at today's speed, never a failed one. + """ + n = derive_process_count(config) + if n < 2: + logger.info("symbol extraction: in-process (extract_processes=1; no pool built)") + return ExtractionPool(n_processes=0) + if not _run_preflight_probe(): + # _run_preflight_probe already logged the WARNING naming the cause. + return ExtractionPool(n_processes=0) + logger.info("symbol extraction: %d process(es) (spawn); pool preflight ok", n) + return ExtractionPool(n_processes=n) diff --git a/indexer/job.py b/indexer/job.py index c5c5c6f..16e8d23 100644 --- a/indexer/job.py +++ b/indexer/job.py @@ -141,6 +141,7 @@ from app.embed import EmbeddingCountMismatchError, EmbedFn, get_embedder from indexer.branches import resolve_branches from indexer.chunk_store import write_chunks_batch +from indexer.extract_pool import ExtractionPool, build_extraction_pool from indexer.fetch import ( REQUIRED_FREE_BYTES, assert_disk_headroom, @@ -165,7 +166,6 @@ reconcile_removed_repos, reconcile_retired_branches, ) -from indexer.symbols import extract_file from indexer.timing import PhaseTimer, install_timer, now, reset_timer logger = logging.getLogger("indexer.job") @@ -309,19 +309,23 @@ def run( reconcile_retired_fn: Callable[..., ReconcileCounts] = reconcile_retired_branches, reconcile_removed_fn: Callable[..., list[str]] = reconcile_removed_repos, shas_fn: Callable[..., ContentShaSets] = read_indexed_shas, + extraction_pool: ExtractionPool | None = None, ) -> int: """Index every configured repo and return a process exit code (0 = all ok). Boundaries are injectable for tests: ``workspace_client`` (secret + config read), ``http_client`` (GitHub HTTP), ``engine`` (DB), ``index_fn`` (the store), ``embed_fn`` (semantic chunking), ``config_loader`` (so orchestration - tests need no SDK fake), and ``reconcile_retired_fn`` / ``reconcile_removed_fn`` + tests need no SDK fake), ``reconcile_retired_fn`` / ``reconcile_removed_fn`` (the storage primitives the post-fan-out checkpoint below invokes -- mirroring - ``index_fn``'s injection so reconciliation tests need no real Postgres either). - ``cfg`` defaults to the process-cached :func:`app.config.get_settings`, then any - ``semantic:`` fields in config.yaml are overlaid onto it before the worker clamp - and embedder build (config.yaml > env > default -- config.yaml is the job's - semantic-config surface; see :class:`indexer.repo_config.SemanticOverrides`). + ``index_fn``'s injection so reconciliation tests need no real Postgres either), + and ``extraction_pool`` (issue #108's shared symbol/edge-extraction pool -- + mirroring ``owns_engine``/``owns_http``, a caller-supplied pool is never + shut down here). ``cfg`` defaults to the process-cached + :func:`app.config.get_settings`, then any ``semantic:`` fields in config.yaml + are overlaid onto it before the worker clamp and embedder build (config.yaml > + env > default -- config.yaml is the job's semantic-config surface; see + :class:`indexer.repo_config.SemanticOverrides`). ``resolve_repos`` is deliberately **not** injectable: the existing ``httpx.MockTransport`` seam already lets tests drive enumeration outcomes @@ -452,6 +456,10 @@ def run( reconciliation_failed = False reconcile_skip_reason = "" progress = ReconcileProgress() + # Computed BEFORE the try below (from the parameter, not the built value) so + # it stays defined even if disk_usage or pool construction never runs -- the + # finally block reads it unconditionally. + owns_pool = extraction_pool is None try: # Inside the try for the same reason as the stamp read below: shutil # .disk_usage raises if tmp is missing or unmounted, and above the try @@ -467,6 +475,27 @@ def run( REQUIRED_FREE_BYTES / 1e9, ) + # Built once per run, shared across every repo worker below (issue + # #108). Wrapped in a bare `except Exception`, not a type list: + # NOTHING about this pool may ever fail the run -- build_extraction_pool + # already degrades its own internal failures (a bad preflight probe, an + # unavailable sandbox) to a WARNING plus an in-process pool, so this + # guard exists only for a truly unexpected failure in that derivation + # itself. `run()` has no handler around the disk-headroom block above + # either, so an uncaught escape here would propagate out of `run()` and + # fail the wheel task after the `finally` below -- this is what prevents + # that. + if extraction_pool is None: + try: + extraction_pool = build_extraction_pool(config) + except Exception: + logger.warning( + "failed to build the shared extraction pool; indexing will " + "run in-process for this run", + exc_info=True, + ) + extraction_pool = ExtractionPool(n_processes=0) + # One batched read for the whole run, BEFORE fan-out. Inside the try so a # failure here still closes the http client and disposes the engine. stamps = _read_stamps(engine, entries) @@ -486,6 +515,7 @@ def run( embed_fn=embed_fn, stamps=stamps, shas_fn=shas_fn, + extraction_pool=extraction_pool, ): entry for entry in entries } @@ -564,6 +594,13 @@ def run( entries=entries, ) finally: + # Shut down BEFORE the http client/engine below, and only ours: an + # injected pool is the caller's to manage (mirrors owns_engine/owns_http). + # By this point the ThreadPoolExecutor `with` block above has already + # exited (every repo worker joined), so no branch holds a future against + # this pool -- the condition that makes shutdown(wait=True) safe. + if owns_pool and extraction_pool is not None: + extraction_pool.shutdown() if owns_http: http_client.close() if owns_engine: @@ -942,6 +979,7 @@ def _index_one( embed_fn: EmbedFn | None, stamps: dict[tuple[str, str], tuple[str | None, int | None]], shas_fn: Callable[..., ContentShaSets], + extraction_pool: ExtractionPool, ) -> RepoOutcome: """Run the full fetch -> parse -> symbols -> store pipeline for every branch of one repo. @@ -968,6 +1006,7 @@ def _index_one( embed_fn=embed_fn, stamps=stamps, shas_fn=shas_fn, + extraction_pool=extraction_pool, ) finally: _repo_ctx.reset(token) @@ -984,6 +1023,7 @@ def _index_one_inner( embed_fn: EmbedFn | None, stamps: dict[tuple[str, str], tuple[str | None, int | None]], shas_fn: Callable[..., ContentShaSets], + extraction_pool: ExtractionPool, ) -> RepoOutcome: """The body of :func:`_index_one`, run with the repo log context already set. @@ -1045,6 +1085,7 @@ def _index_one_inner( started=started, max_chunks_per_repo=max_chunks_per_repo, shas_fn=shas_fn, + extraction_pool=extraction_pool, ) for branch in resolution.branches ] @@ -1113,6 +1154,7 @@ def _index_one_branch( started: float, max_chunks_per_repo: int, shas_fn: Callable[..., ContentShaSets], + extraction_pool: ExtractionPool, ) -> BranchOutcome: """Fetch, parse, and store ONE branch. Never raises -- every failure is classified. @@ -1152,6 +1194,15 @@ def _index_one_branch( ``_index_one``'s ``_repo_ctx`` reset is: ``ThreadPoolExecutor`` reuses worker threads without resetting their context, so a leaked timer would attribute this branch's sweep to the next task on this thread. + + ``extraction_pool`` (issue #108) is the run-shared, ``spawn``-based process + pool ``indexer.extract_pool`` extraction runs through; both ``items`` + expressions below call ``extraction_pool.stream(...)`` rather than + ``extract_file`` directly. A worker death raises + ``indexer.extract_pool.ExtractionPoolError`` out of that stream -- caught by + THIS function's existing broad ``except Exception`` below like any other + extraction failure, with no new ``except`` branch: legibility comes from + ``ExtractionPoolError``'s own message, not from special-casing it here. """ # NOT the `started` parameter: that one is set once per REPO and passed # unchanged to every branch, so reusing it would make branch 2+ of a @@ -1273,10 +1324,12 @@ def _index_one_branch( precompute_failed = True finally: timer.add("embed", timer.clock() - t0) - items = ((pf, extract_file(pf)) for pf in files) + items = extraction_pool.stream(files) else: # Lazy generator: files stream through the open transaction (bounded memory). - items = ((pf, extract_file(pf)) for pf in iter_tar_source_files(tar_path)) + # extraction_pool.stream() is itself pull-driven and single-consumer, so this + # stays lazy end to end -- no full materialization on the semantic-off path. + items = extraction_pool.stream(iter_tar_source_files(tar_path)) # Parse time is INSIDE the db window (items are produced lazily as # index_repo consumes them), so `db` subtracts only the parse accrued diff --git a/indexer/repo_config.py b/indexer/repo_config.py index d863771..38cacf5 100644 --- a/indexer/repo_config.py +++ b/indexer/repo_config.py @@ -264,9 +264,11 @@ class RepoConfig(BaseModel): numbers 5x but deliberately left the default alone; #109 re-derives it.) **Returns at the ceiling are sublinear.** Symbol extraction does not - parallelise (measured at 0.95x on 4 threads), so Amdahl's law caps the - speedup well below 8x while the disk cost stays a hard linear 4 GB. Raise - it only knowing that trade. + parallelise across THREADS (measured at 0.95x on 4 threads -- the tree walk + is GIL-serialized) -- which is exactly why extraction now runs in a shared + process pool instead (``extract_processes``, below); that pool is decoupled + from ``index_concurrency`` entirely. Raise ``index_concurrency`` only + knowing you are buying disk-bound repo fan-out, not extraction throughput. When semantic indexing is on, the effective worker count is clamped to 2 by :func:`effective_workers`. That clamp is a **memory** bound, not a CPU one: @@ -286,11 +288,25 @@ class RepoConfig(BaseModel): the job resolves ``entry.semantic_max_chunks or cfg.semantic_max_chunks_per_repo`` and the ``semantic:`` block only supplies the second operand -- see :class:`SemanticOverrides`. + + ``extract_processes`` (#108) sizes the shared, ``spawn``-based process pool + ``indexer.extract_pool`` runs symbol/edge extraction through -- one pool per + run, shared across every ``index_concurrency`` repo worker, so a single + giant repo's files parse on every available core rather than one thread's + worth. It is a CPU knob, entirely independent of ``index_concurrency``'s + disk bound above. ``None`` (the default) derives from the runtime: + affinity/cgroup-aware CPU count, clamped to 8 -- matching every other + parallelism knob in this repo (``index_concurrency``, + ``SemanticOverrides.embedding_concurrency``). Setting this to ``1`` restores + fully serial, in-process extraction and spawns no process pool at all -- + the rollback switch, exercised the same way #107's + ``embedding_concurrency: 1`` is. """ version: Literal[1] connections: list[Connection] = Field(min_length=1) index_concurrency: int = Field(default=4, ge=1, le=8) + extract_processes: int | None = Field(default=None, ge=1, le=8) # Per-repo override of Settings.semantic_max_chunks_per_repo (app/config.py, default # 8000), keyed by repo. Absent (the default) -> no repo gets an override, and diff --git a/scripts/measure_extraction_pool.py b/scripts/measure_extraction_pool.py new file mode 100644 index 0000000..1534778 --- /dev/null +++ b/scripts/measure_extraction_pool.py @@ -0,0 +1,128 @@ +"""Measure symbol/edge extraction pool scaling against a REAL tarball (#108, AC1). + +Offline companion to :mod:`indexer.extract_pool`. Drives the exact production call +site :mod:`indexer.job` uses: ``indexer.ingest.iter_tar_source_files(tarball)`` -> +``ExtractionPool.stream(...)`` -- never a batch-submit harness (``executor.map`` +over a pre-built batch list) and never an on-disk walk. Since #106 the production +file source is a single serial gzip-decompress + decode + filter pass over one +open ``TarFile``, run once here and shared between the serial and pooled arms, so +this script also reports the ingest/extract split the pool cannot touch (see +``docs/runbooks/indexing-parallelism.md`` §3's "ingestion stays serial" note). + +**The tarball must be real** -- a GitHub codeload archive (``curl -L +https://codeload.github.com/OWNER/REPO/tar.gz/HEAD -o repo.tar.gz`` or +``indexer.fetch.download_tarball`` by hand), not a synthetic fixture. A +site-packages-heavy Python tree (e.g. this repo's own ``.venv/``, tarred up) is a +reasonable large-repo proxy: ``tar czf repo.tar.gz --transform 's,^,repo-abc1234/,' .`` + +Usage: ``uv run python scripts/measure_extraction_pool.py --tarball repo.tar.gz +[--processes 2,4,8] [--batch-bytes 2000000] [--repeat 1]`` +""" + +from __future__ import annotations + +import argparse +import time +from collections.abc import Sequence +from pathlib import Path + +from indexer.extract_pool import _BATCH_BYTES, ExtractionPool +from indexer.ingest import iter_tar_source_files +from indexer.languages import FileExtraction, ParsedFile +from indexer.symbols import extract_file + + +def _load_files(tarball: Path) -> list[ParsedFile]: + """One real, serial pass over the tarball -- the Amdahl floor (§4.7.1) this + pool cannot parallelize. Timed separately by the caller.""" + return list(iter_tar_source_files(tarball)) + + +def _measure_serial( + files: list[ParsedFile], +) -> tuple[float, list[tuple[ParsedFile, FileExtraction]]]: + start = time.perf_counter() + results = [(pf, extract_file(pf)) for pf in files] + return time.perf_counter() - start, results + + +def _measure_pooled( + files: list[ParsedFile], *, n_processes: int, batch_bytes: int +) -> tuple[float, list[tuple[ParsedFile, FileExtraction]]]: + pool = ExtractionPool(n_processes=n_processes, batch_bytes=batch_bytes) + try: + start = time.perf_counter() + results = list(pool.stream(files)) + elapsed = time.perf_counter() - start + finally: + pool.shutdown() + return elapsed, results + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--tarball", required=True, type=Path, help="a real downloaded repo tarball (.tar.gz)" + ) + parser.add_argument( + "--processes", default="2,4,8", help="comma-separated worker-process counts to measure" + ) + parser.add_argument("--batch-bytes", type=int, default=_BATCH_BYTES) + parser.add_argument("--repeat", type=int, default=1, help="repetitions per data point") + args = parser.parse_args(argv) + process_counts = [int(p) for p in args.processes.split(",")] + + print(f"loading {args.tarball} ...") + ingest_start = time.perf_counter() + files = _load_files(args.tarball) + ingest_elapsed = time.perf_counter() - ingest_start + qualifying = [pf for pf in files if pf.lang is not None] + total_bytes = sum(len(pf.content) for pf in files) + qualifying_bytes = sum(len(pf.content) for pf in qualifying) + print( + f"{len(files)} indexable files / {total_bytes / 1e6:.1f} MB, " + f"{len(qualifying)} qualify for parsing / {qualifying_bytes / 1e6:.1f} MB " + f"(ingest: {ingest_elapsed:.2f}s -- serial in the parent, both arms below)" + ) + + print(f"\n{'path':>14s} {'extract_s':>10s} {'total_s':>10s} {'speedup':>8s} {'note':>9s}") + serial_extract = None + serial_results: list[tuple[ParsedFile, FileExtraction]] = [] + for _ in range(args.repeat): + elapsed, serial_results = _measure_serial(files) + if serial_extract is None or elapsed < serial_extract: + serial_extract = elapsed + assert serial_extract is not None + serial_total = ingest_elapsed + serial_extract + print(f"{'serial':>14s} {serial_extract:>10.3f} {serial_total:>10.3f} {1.0:>7.2f}x") + + for n in process_counts: + best_elapsed = None + identical = True + for _ in range(args.repeat): + elapsed, pooled_results = _measure_pooled( + files, n_processes=n, batch_bytes=args.batch_bytes + ) + identical = identical and pooled_results == serial_results + if best_elapsed is None or elapsed < best_elapsed: + best_elapsed = elapsed + assert best_elapsed is not None + total = ingest_elapsed + best_elapsed + speedup = serial_total / total if total else float("inf") + note = "identical" if identical else "MISMATCH" + print( + f"{f'pool x{n}':>14s} {best_elapsed:>10.3f} {total:>10.3f} " + f"{speedup:>7.2f}x {note:>9s}" + ) + + print( + f"\ningest (serial, shared by both arms): {ingest_elapsed:.2f}s of " + f"{serial_total:.2f}s serial total ({100 * ingest_elapsed / serial_total:.0f}%) -- " + "the floor AC1's 'combined' speedup above cannot cross, however many " + "processes extract_processes uses." + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_extract_pool.py b/tests/unit/test_extract_pool.py new file mode 100644 index 0000000..baa3cfb --- /dev/null +++ b/tests/unit/test_extract_pool.py @@ -0,0 +1,626 @@ +"""Unit tests for indexer.extract_pool: the shared process-pool extraction (#108). + +T1-T8 below mirror the execution plan's numbering. T9/T10 (the existing +``test_job.py`` suite's kill-switch pin and phase-timing attribution) live in +``tests/unit/test_job.py`` -- this file covers ``indexer.extract_pool`` in +isolation. +""" + +from __future__ import annotations + +import logging +import threading +import time +from concurrent.futures import Future +from typing import Any + +import pytest + +from indexer import extract_pool +from indexer.extract_pool import ( + _MAX_BATCH_FILES, + MAX_POOL_REBUILDS, + ExtractionPool, + ExtractionPoolError, + _needs_parse, + build_extraction_pool, + derive_process_count, +) +from indexer.languages import EXT_TO_LANG, SYMBOL_KINDS, FileExtraction, ParsedFile +from indexer.repo_config import RepoConfig +from indexer.symbols import extract_file + +# --- fixtures ----------------------------------------------------------------- + +# One real per-language snippet each (reusing test_symbols.py's shapes), so T1's +# parity check exercises every SYMBOL_KINDS language through a real spawn pool. +_LANG_SOURCES: dict[str, str] = { + "python": "class C:\n def m(self):\n pass\ndef top():\n pass\n", + "javascript": "class C {\n m() {}\n}\nfunction top() {}\n", + "typescript": "interface I {}\nclass C {\n m() {}\n}\nfunction top() {}\n", + "tsx": "function top() {}\n", + "go": "package main\nfunc top() {}\nfunc (r R) m() {}\ntype T int\n", + "java": "class C {\n void m() {}\n}\ninterface I {}\n", + "rust": "fn top() {}\nstruct S;\nenum E {}\ntrait T {}\n", +} + + +def _fixture_corpus() -> list[ParsedFile]: + files = [ + ParsedFile(path=f"src/f.{lang}", lang=lang, size=len(content), content=content) + for lang, content in _LANG_SOURCES.items() + ] + files.append(ParsedFile(path="unknown.dat", lang=None, size=3, content="hi\n")) + files.append(ParsedFile(path="doc.md", lang="markdown", size=5, content="# hi\n")) + files.append(ParsedFile(path="empty.py", lang="python", size=0, content="")) + files.append( + ParsedFile( + path="nested.py", + lang="python", + size=0, + content=( + "class Outer:\n" + " def method(self):\n" + " return helper()\n" + "def helper():\n" + " return 1\n" + ), + ) + ) + return files + + +def _python_files(n: int, *, size: int = 200) -> list[ParsedFile]: + padding = "x" * max(size - 20, 1) + return [ + ParsedFile( + path=f"pkg/mod{i}.py", + lang="python", + size=size, + content=f"def f{i}():\n return '{padding}'\n", + ) + for i in range(n) + ] + + +def _config(*, extract_processes: int | None = None) -> RepoConfig: + doc: dict[str, Any] = {"version": 1, "connections": [{"type": "github", "users": ["u"]}]} + if extract_processes is not None: + doc["extract_processes"] = extract_processes + return RepoConfig.model_validate(doc) + + +# Module-level (picklable under spawn) preflight-probe targets for T6a/T6b -- +# see _run_preflight_probe's docstring: a test override must be a module-level +# function, never a closure or lambda, or spawn's re-import cannot find it. +def _crashing_probe_target(queue: Any) -> None: + raise RuntimeError("simulated preflight crash") + + +def _hanging_probe_target(queue: Any) -> None: + time.sleep(3600) + + +# --- T1: parity, real spawn pool ---------------------------------------------- + + +@pytest.mark.unit +def test_stream_parity_against_extract_file_real_spawn_pool() -> None: + """T1 (AC2): real spawn processes, real pickling -- the risk the issue names. + + ``max_workers=2`` is pinned deliberately: each spawned child re-imports + ``tree_sitter_language_pack``, so worker count drives this test's cost more + than corpus size does. A tiny ``batch_bytes`` forces the fixture corpus + (well under the real 2 MB default) across >=3 batches, so submit/collect + round-trips more than once. + """ + files = _fixture_corpus() + pool = ExtractionPool(n_processes=2, batch_bytes=20) + try: + got = list(pool.stream(files)) + finally: + pool.shutdown() + + expected = [(pf, extract_file(pf)) for pf in files] + assert got == expected + + +@pytest.mark.unit +def test_build_extraction_pool_real_spawn_end_to_end() -> None: + """Sanity: the full build_extraction_pool -> preflight -> stream wiring, + through a real spawn pool (not the lower-level ExtractionPool constructor + T1 uses directly).""" + pool = build_extraction_pool(_config(extract_processes=2)) + try: + assert pool._executor is not None + files = _python_files(5) + got = list(pool.stream(files)) + finally: + pool.shutdown() + + expected = [(pf, extract_file(pf)) for pf in files] + assert got == expected + + +# --- T2: routing-predicate agreement ------------------------------------------ + + +@pytest.mark.unit +@pytest.mark.parametrize("lang", [*sorted(EXT_TO_LANG.values()), None, "bogus"]) +def test_needs_parse_agrees_with_extract_file_short_circuit(lang: str | None) -> None: + pf = ParsedFile(path="x", lang=lang, size=0, content="whatever\n") + expect_parse = lang is not None and lang in SYMBOL_KINDS + assert _needs_parse(pf) is expect_parse + if not expect_parse: + assert extract_file(pf) == FileExtraction(symbols=[], edges=[]) + + +# --- T3: bounded look-ahead ---------------------------------------------------- + + +class _CountingSource: + """Wraps a file list, counting how many items have actually been pulled.""" + + def __init__(self, files: list[ParsedFile]) -> None: + self._it = iter(files) + self.pulled = 0 + + def __iter__(self) -> "_CountingSource": + return self + + def __next__(self) -> ParsedFile: + pf = next(self._it) + self.pulled += 1 + return pf + + +class _ImmediateFuture: + def __init__(self, result: list[FileExtraction]) -> None: + self._result = result + + def result(self) -> list[FileExtraction]: + return self._result + + def cancel(self) -> bool: + return False + + +class _ImmediateExecutor: + """Resolves every submitted batch synchronously -- still lets ExtractionPool's + own in_flight bookkeeping (not completion timing) drive the look-ahead bound.""" + + def submit(self, fn: Any, *args: Any) -> _ImmediateFuture: + return _ImmediateFuture(fn(*args)) + + +@pytest.mark.unit +def test_bounded_look_ahead() -> None: + """T3: the issue's explicit no-full-materialization rule. Each file alone + exceeds the tiny batch_bytes, so one file == one batch, and max_in_flight + (2 * n_processes) directly bounds how many files are pulled ahead.""" + total = 20 + files = _python_files(total, size=200) + source = _CountingSource(files) + + pool = ExtractionPool(n_processes=0, batch_bytes=50) + pool._n_processes = 2 + pool._executor = _ImmediateExecutor() # type: ignore[assignment] + max_in_flight = 2 * pool._n_processes + + gen = pool.stream(source) # type: ignore[arg-type] + next(gen) + assert source.pulled <= max_in_flight + assert source.pulled < total + + for _ in gen: + assert source.pulled <= total + assert source.pulled == total + + +@pytest.mark.unit +def test_bounded_look_ahead_for_non_qualifying_files() -> None: + """Regression: non-qualifying (locally-answered) files never enter a batch, + so a look-ahead bound keyed on in-flight BATCH count alone never throttles a + corpus dominated by them -- the whole source would be pulled in one shot. + The bound must cover pending local content bytes too.""" + total = 500 + files = [ + ParsedFile(path=f"doc{i}.md", lang=None, size=2000, content="x" * 2000) + for i in range(total) + ] + source = _CountingSource(files) + + # batch_bytes deliberately tiny: max_pending_bytes (4 * batch_bytes) must be + # well under the corpus's total bytes (500 * 2000 = 1,000,000) for this test + # to actually exercise the throttle rather than trivially fitting everything + # under budget. + pool = ExtractionPool(n_processes=0, batch_bytes=1000) + pool._n_processes = 2 + pool._executor = _ImmediateExecutor() # type: ignore[assignment] + + gen = pool.stream(source) # type: ignore[arg-type] + next(gen) + assert source.pulled < total # bounded -- NOT the whole corpus in one shot + + for _ in gen: + pass + assert source.pulled == total + + +class _RecordingExecutor: + """Resolves synchronously like _ImmediateExecutor, but also records each + submitted batch's file count -- for pinning the per-batch file-count cap.""" + + def __init__(self) -> None: + self.batch_sizes: list[int] = [] + + def submit(self, fn: Any, *args: Any) -> _ImmediateFuture: + self.batch_sizes.append(len(args[0])) + return _ImmediateFuture(fn(*args)) + + +@pytest.mark.unit +def test_batch_is_capped_by_file_count_even_with_huge_batch_bytes() -> None: + """Regression: a corpus of near-zero-byte qualifying files barely moves + batch_bytes, so without a per-batch FILE-COUNT cap a single batch would grow + to the entire look-ahead window (thousands of files, one oversized pickle on + one idle worker) before ever tripping the byte bound.""" + total = _MAX_BATCH_FILES * 2 + 5 + files = _python_files(total, size=1) # near-zero content per file + executor = _RecordingExecutor() + + pool = ExtractionPool(n_processes=0, batch_bytes=10_000_000) # never trips + pool._n_processes = 2 + pool._executor = executor # type: ignore[assignment] + + got = list(pool.stream(files)) + assert got == [(pf, extract_file(pf)) for pf in files] + assert executor.batch_sizes # at least one batch was actually submitted + assert max(executor.batch_sizes) <= _MAX_BATCH_FILES + + +# --- T4: order preservation under out-of-order completion --------------------- + + +class _ReverseOrderExecutor: + """Batch N's future resolves BEFORE batch N-1's -- proves stream() yields in + submission (source) order, not completion order.""" + + def __init__(self, total_batches: int) -> None: + self._total = total_batches + self._n = 0 + + def submit(self, fn: Any, *args: Any) -> "Future[list[FileExtraction]]": + self._n += 1 + idx = self._n + fut: "Future[list[FileExtraction]]" = Future() + + def _worker() -> None: + time.sleep(0.01 * (self._total - idx + 1)) + fut.set_result(fn(*args)) + + threading.Thread(target=_worker, daemon=True).start() + return fut + + +@pytest.mark.unit +def test_order_preservation_under_out_of_order_completion() -> None: + """T4 (D6): later-submitted batches resolving first must not reorder output.""" + files = _python_files(6, size=200) # tiny batch_bytes -> one batch per file + pool = ExtractionPool(n_processes=0, batch_bytes=50) + pool._n_processes = 2 + pool._executor = _ReverseOrderExecutor(len(files)) # type: ignore[assignment] + + got = list(pool.stream(files)) + expected = [(pf, extract_file(pf)) for pf in files] + assert got == expected + + +# --- T5: BrokenProcessPool isolation, two concurrent streams ------------------ + + +class _BrokenFuture: + def __init__(self, barrier: threading.Barrier | None = None) -> None: + self._barrier = barrier + + def result(self) -> list[FileExtraction]: + if self._barrier is not None: + # Rendezvous so two concurrent callers genuinely observe the SAME + # (pre-rebuild) executor before either raises -- without this, one + # thread can race ahead, rebuild, and hand the second thread an + # already-healthy pool, which would defeat the point of this test. + self._barrier.wait(timeout=10.0) + raise extract_pool.BrokenProcessPool("worker died") + + def cancel(self) -> bool: + return False + + +class _AlwaysBreakingExecutor: + def __init__(self, barrier: threading.Barrier | None = None) -> None: + self._barrier = barrier + + def submit(self, fn: Any, *args: Any) -> _BrokenFuture: + return _BrokenFuture(self._barrier) + + def shutdown(self, wait: bool = True) -> None: + pass + + +class _WorkingExecutor: + def submit(self, fn: Any, *args: Any) -> "Future[list[FileExtraction]]": + fut: "Future[list[FileExtraction]]" = Future() + fut.set_result(fn(*args)) + return fut + + def shutdown(self, wait: bool = True) -> None: + pass + + +@pytest.mark.unit +def test_broken_process_pool_isolation_two_concurrent_streams() -> None: + """T5 (AC3): the two-thread shape is the point -- a single-stream test would + pass against a racy rebuild.""" + files = _python_files(3, size=50) + pool = ExtractionPool(n_processes=0, batch_bytes=10) + pool._n_processes = 2 + barrier = threading.Barrier(2) + pool._executor = _AlwaysBreakingExecutor(barrier=barrier) # type: ignore[assignment] + + rebuild_calls = {"n": 0} + + def _new_executor_stub() -> Any: + rebuild_calls["n"] += 1 + # The FIRST rebuild recovers onto a healthy executor (proves "a + # subsequent stream succeeds on the new generation"); every rebuild + # after that keeps breaking, driving toward the rebuild budget. + return _WorkingExecutor() if rebuild_calls["n"] == 1 else _AlwaysBreakingExecutor() + + pool._new_executor = _new_executor_stub # type: ignore[method-assign] + + errors: list[ExtractionPoolError] = [] + lock = threading.Lock() + + def _drive() -> None: + try: + list(pool.stream(files)) + except ExtractionPoolError as exc: + with lock: + errors.append(exc) + + t1 = threading.Thread(target=_drive) + t2 = threading.Thread(target=_drive) + t1.start() + t2.start() + t1.join() + t2.join() + + assert len(errors) == 2 # both concurrent callers saw the failure + assert isinstance(errors[0].__cause__, extract_pool.BrokenProcessPool) + assert "extract_processes: 1" in str(errors[0]) # names the rollback + assert pool._rebuilds == 1 # exactly one rebuild for N concurrent breaks + assert pool._generation == 1 + assert not pool._latched + + # A subsequent stream succeeds on the new (healthy) generation. + got = list(pool.stream(files)) + assert got == [(pf, extract_file(pf)) for pf in files] + + # Simulate a later worker death under the SAME (still healthy) generation, + # driving repeated rebuilds until the budget is exhausted and it latches. + pool._executor = _AlwaysBreakingExecutor() # type: ignore[assignment] + for _ in range(MAX_POOL_REBUILDS - 1): + with pytest.raises(ExtractionPoolError): + list(pool.stream(files)) + assert pool._rebuilds == MAX_POOL_REBUILDS + assert not pool._latched + + with pytest.raises(ExtractionPoolError): + list(pool.stream(files)) + assert pool._latched + assert pool._executor is None + + # Latched: still yields correct results, now in-process. + got_latched = list(pool.stream(files)) + assert got_latched == [(pf, extract_file(pf)) for pf in files] + + +class _BreaksOnSubmitExecutor: + """Unlike _AlwaysBreakingExecutor (which fails at .result()), this fails + inside submit() itself -- the shape a real ProcessPoolExecutor takes once a + worker has already died before the NEXT batch is even dispatched.""" + + def submit(self, fn: Any, *args: Any) -> Any: + raise extract_pool.BrokenProcessPool("worker already dead at submit time") + + def shutdown(self, wait: bool = True) -> None: + pass + + +@pytest.mark.unit +def test_broken_process_pool_raised_from_submit_is_handled() -> None: + """Regression: a real ProcessPoolExecutor can raise BrokenProcessPool from + submit() itself, not only from a future's result() -- verified against a + real ProcessPoolExecutor in review. The supervisor must still record the + break (rebuild-or-latch) rather than let the raw exception escape + unclassified and leave (_rebuilds, _generation, _latched) stale.""" + files = _python_files(3, size=50) + pool = ExtractionPool(n_processes=0, batch_bytes=10) + pool._n_processes = 2 + pool._executor = _BreaksOnSubmitExecutor() # type: ignore[assignment] + + with pytest.raises(ExtractionPoolError) as excinfo: + list(pool.stream(files)) + + assert isinstance(excinfo.value.__cause__, extract_pool.BrokenProcessPool) + assert pool._rebuilds == 1 + assert pool._generation == 1 + assert not pool._latched + + +# --- T6a/T6b: preflight probe ------------------------------------------------- + + +@pytest.mark.unit +def test_preflight_fails_fast_when_ipc_queue_cannot_be_created( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """T6a: the sem_open/OSError half of "probe child exits nonzero (or + ctx.Queue() raises OSError)". Fails before any child is even started.""" + + class _BrokenContext: + def Queue(self) -> Any: + raise OSError("sem_open: No space left on device") + + monkeypatch.setattr(extract_pool.multiprocessing, "get_context", lambda name: _BrokenContext()) + + with caplog.at_level(logging.WARNING, logger="indexer.extract_pool"): + ok = extract_pool._run_preflight_probe() + + assert ok is False + assert "could not create an IPC queue" in caplog.text + + +@pytest.mark.unit +def test_preflight_fails_fast_when_probe_child_crashes(caplog: pytest.LogCaptureFixture) -> None: + """T6a: a real spawned child that exits nonzero is detected -- and the WARNING + names the cause.""" + with caplog.at_level(logging.WARNING, logger="indexer.extract_pool"): + ok = extract_pool._run_preflight_probe(target=_crashing_probe_target, timeout=10.0) + + assert ok is False + assert "exited with code" in caplog.text + + +@pytest.mark.unit +def test_preflight_hang_is_killed_and_latches(caplog: pytest.LogCaptureFixture) -> None: + """T6b: THE defect this plan was blocked on. A hung probe child must be + killed within the injected timeout, never awaited to completion, and the + test process itself must not hang.""" + start = time.monotonic() + with caplog.at_level(logging.WARNING, logger="indexer.extract_pool"): + ok = extract_pool._run_preflight_probe(target=_hanging_probe_target, timeout=0.5) + elapsed = time.monotonic() - start + + assert ok is False + assert elapsed < 30.0 # nowhere near the child's 3600s sleep + assert "did not respond within" in caplog.text + assert "killed" in caplog.text + + +@pytest.mark.unit +def test_preflight_failure_latches_build_extraction_pool_to_in_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """D8: build_extraction_pool degrades to a correct in-process pool, and no + ProcessPoolExecutor is ever constructed, when the preflight probe fails.""" + monkeypatch.setattr(extract_pool, "_run_preflight_probe", lambda: False) + + pool = build_extraction_pool(_config(extract_processes=4)) + + assert pool._executor is None + files = _python_files(3) + assert list(pool.stream(files)) == [(pf, extract_file(pf)) for pf in files] + + +# --- T7: sizing derivation ------------------------------------------------------ + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("content", "expected"), + [ + ("200000 100000\n", 2), + ("100000 100000\n", 1), + ("max 100000\n", None), + ("bogus 100000\n", None), + ("200000 0\n", None), + ("200000\n", None), + ("\n", None), + ], +) +def test_cgroup_cpu_quota_parsing( + monkeypatch: pytest.MonkeyPatch, content: str, expected: int | None +) -> None: + monkeypatch.setattr(extract_pool.Path, "read_text", lambda self: content) + assert extract_pool._cgroup_cpu_quota() == expected + + +@pytest.mark.unit +def test_cgroup_cpu_quota_unreadable_file_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: + def _raise(self: Any) -> str: + raise OSError("no such file") + + monkeypatch.setattr(extract_pool.Path, "read_text", _raise) + assert extract_pool._cgroup_cpu_quota() is None + + +@pytest.mark.unit +def test_available_cpus_prefers_affinity_over_cpu_count(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(extract_pool.os, "sched_getaffinity", lambda pid: {0, 1, 2}, raising=False) + monkeypatch.setattr(extract_pool.os, "cpu_count", lambda: 64) + monkeypatch.setattr(extract_pool, "_cgroup_cpu_quota", lambda: None) + assert extract_pool._available_cpus() == 3 + + +@pytest.mark.unit +def test_available_cpus_falls_back_to_cpu_count_without_affinity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delattr(extract_pool.os, "sched_getaffinity", raising=False) + monkeypatch.setattr(extract_pool.os, "cpu_count", lambda: 5) + monkeypatch.setattr(extract_pool, "_cgroup_cpu_quota", lambda: None) + assert extract_pool._available_cpus() == 5 + + +@pytest.mark.unit +def test_available_cpus_clamped_by_cgroup_quota(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + extract_pool.os, "sched_getaffinity", lambda pid: set(range(16)), raising=False + ) + monkeypatch.setattr(extract_pool, "_cgroup_cpu_quota", lambda: 2) + assert extract_pool._available_cpus() == 2 + + +@pytest.mark.unit +def test_available_cpus_floors_at_one(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(extract_pool.os, "sched_getaffinity", lambda pid: set(), raising=False) + monkeypatch.setattr(extract_pool, "_cgroup_cpu_quota", lambda: None) + assert extract_pool._available_cpus() == 1 + + +@pytest.mark.unit +def test_derive_process_count_explicit_value_wins(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(extract_pool, "_available_cpus", lambda: 1) # would derive to 1 if used + assert derive_process_count(_config(extract_processes=6)) == 6 + + +@pytest.mark.unit +def test_derive_process_count_kill_switch_is_explicit_one() -> None: + assert derive_process_count(_config(extract_processes=1)) == 1 + + +@pytest.mark.unit +def test_derive_process_count_derives_and_clamps_to_eight(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(extract_pool, "_available_cpus", lambda: 64) + assert derive_process_count(_config()) == 8 + + +@pytest.mark.unit +def test_derive_process_count_derives_below_ceiling(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(extract_pool, "_available_cpus", lambda: 3) + assert derive_process_count(_config()) == 3 + + +# --- T8: extract_processes: 1 creates no pool ---------------------------------- + + +@pytest.mark.unit +def test_extract_processes_one_creates_no_pool() -> None: + """T8: the kill switch. No subprocess is ever created (no executor object at + all), and results are identical to calling extract_file directly.""" + pool = build_extraction_pool(_config(extract_processes=1)) + assert pool._executor is None + + files = _python_files(3) + assert list(pool.stream(files)) == [(pf, extract_file(pf)) for pf in files] diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 8d1c4a9..8fe93fe 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -25,6 +25,7 @@ import tarfile import threading import time +from concurrent.futures import Future from pathlib import Path from typing import Any, NamedTuple @@ -33,6 +34,7 @@ from app.config import Settings from app.db.models import INDEX_SEMANTICS_VERSION +from indexer.extract_pool import ExtractionPool, ExtractionPoolError from indexer.hashing import content_sha from indexer.job import ( BranchOutcome, @@ -43,7 +45,7 @@ read_github_token, run, ) -from indexer.languages import ExtractedSymbol, IndexCounts, ParsedFile +from indexer.languages import ExtractedSymbol, FileExtraction, IndexCounts, ParsedFile from indexer.repo_config import ConfigError, RepoConfig, load_config from indexer.resolve import RepoEntry from indexer.store import ( @@ -191,6 +193,7 @@ def _repo_meta(full_name: str, **overrides: Any) -> dict[str, Any]: def _config( *, index_concurrency: int | None = None, + extract_processes: int | None = 1, semantic_max_chunks_per_repo: dict[str, int] | None = None, semantic: dict[str, Any] | None = None, **connection: Any, @@ -201,10 +204,21 @@ def _config( job-side semantic-config surface), mirroring the ``semantic_max_chunks_per_repo`` per-repo-map kwarg above -- the two are distinct config surfaces (global INT vs per-repo MAP) and either can be set independently. + + ``extract_processes`` (#108) defaults to ``1`` -- the pool KILL SWITCH -- + deliberately, not to the schema's own ``None`` default. Left at ``None``, every + one of the ~150 ``run()`` tests below would run the preflight probe and spawn + real worker processes on every call, each re-importing + ``tree_sitter_language_pack``: minutes slower and flaky on a constrained CI + runner. Only the pool-specific tests in ``tests/unit/test_extract_pool.py`` + (and the dedicated pool-lifecycle tests here) opt into a real pool by passing + ``extract_processes=`` explicitly. """ doc: dict[str, Any] = {"version": 1, "connections": [{"type": "github", **connection}]} if index_concurrency is not None: doc["index_concurrency"] = index_concurrency + if extract_processes is not None: + doc["extract_processes"] = extract_processes if semantic_max_chunks_per_repo is not None: doc["semantic_max_chunks_per_repo"] = semantic_max_chunks_per_repo if semantic is not None: @@ -395,6 +409,7 @@ def _run( reconcile_retired_fn: Any = _noop_retired_fn, reconcile_removed_fn: Any = _noop_removed_fn, shas_fn: Any = _noop_shas_fn, + extraction_pool: Any = None, ) -> int: """Drive run() with a faked config read but a REAL resolve_repos. @@ -409,6 +424,11 @@ def _run( version-matching, sha-mismatching stamp with semantic indexing on would otherwise reach the REAL ``read_indexed_shas`` against ``_FakeConn``, which cannot answer it (see that fake's docstring). + + ``extraction_pool`` (#108) defaults to ``None``, in which case ``run()`` + builds its own from ``config.extract_processes`` (pinned to ``1``, the kill + switch, by ``_config``'s own default -- see T9). Pass an explicit fake/real + :class:`~indexer.extract_pool.ExtractionPool` to test the pool seam itself. """ wc = _FakeWorkspaceClient("tok") engine = engine if engine is not None else _FakeEngine() @@ -430,6 +450,7 @@ def _run( reconcile_retired_fn=reconcile_retired_fn, reconcile_removed_fn=reconcile_removed_fn, shas_fn=shas_fn, + extraction_pool=extraction_pool, ) @@ -1070,6 +1091,7 @@ def _down(_texts: list[str]) -> list[list[float]]: started=time.monotonic(), max_chunks_per_repo=100, shas_fn=_noop_shas_fn, + extraction_pool=ExtractionPool(n_processes=0), ) assert outcome.status == "indexed" assert outcome.semantic_degraded is True @@ -1576,6 +1598,7 @@ def test_index_one_inner_reports_discovery_complete_for_a_normal_run() -> None: embed_fn=None, stamps={}, shas_fn=_noop_shas_fn, + extraction_pool=ExtractionPool(n_processes=0), ) assert outcome.name == "acme/widgets" assert outcome.discovery_complete is True @@ -1603,6 +1626,7 @@ def test_index_one_inner_reports_discovery_incomplete_when_capped() -> None: embed_fn=None, stamps={}, shas_fn=_noop_shas_fn, + extraction_pool=ExtractionPool(n_processes=0), ) assert outcome.discovery_complete is False assert len(outcome.outcomes) == SOFT_BRANCH_CAP @@ -1631,6 +1655,7 @@ def test_index_one_inner_default_flip_mirror_is_complete() -> None: embed_fn=None, stamps={}, shas_fn=_noop_shas_fn, + extraction_pool=ExtractionPool(n_processes=0), ) assert [o.branch for o in outcome.outcomes] == ["main"] assert outcome.discovery_complete is True @@ -1708,7 +1733,9 @@ def test_two_repos_are_indexed_concurrently_at_concurrency_two() -> None: NOTE this proves CONCURRENCY, not throughput. It would pass identically in a world where fan-out delivers zero speedup — symbol extraction itself measured - 0.95x on 4 threads. Throughput is measured on the first production run, via + 0.95x on 4 THREADS, which is exactly why extraction now runs in a shared + process pool instead (#108) rather than leaning on this thread-level + fan-out for it. Throughput is measured on the first production run, via the duration log lines asserted below. """ idx = _BarrierIndex(parties=2) @@ -1939,6 +1966,7 @@ def test_repo_context_is_reset_even_when_the_repo_fails() -> None: embed_fn=None, stamps={}, shas_fn=_noop_shas_fn, + extraction_pool=ExtractionPool(n_processes=0), ) assert _repo_ctx.get() == "-" @@ -2933,17 +2961,24 @@ def test_parse_time_is_excluded_from_db_time( ) -> None: """T5: items are produced lazily INSIDE index_repo's transaction, so a naive "time index_fn" would report parse and db fused. _timed_items charges only the - production of each item to parse, leaving the DML between them to db.""" - import indexer.job as job + production of each item to parse, leaving the DML between them to db. + + Patches ``indexer.extract_pool.extract_file`` (#108 moved the in-process + extraction call site there from ``indexer.job``) -- with the kill switch + pinned by ``_config``'s default, ``ExtractionPool.stream()`` degrades to + calling exactly that name per file, so this still measures the real + ``_timed_items`` wrap around the pool's in-process path. + """ + import indexer.extract_pool as extract_pool clock = _install_fake_clock(monkeypatch) - real_extract_file = job.extract_file + real_extract_file = extract_pool.extract_file def _slow_extract(pf: ParsedFile) -> Any: clock.advance(2.0) return real_extract_file(pf) - monkeypatch.setattr(job, "extract_file", _slow_extract) + monkeypatch.setattr(extract_pool, "extract_file", _slow_extract) def _index(conn: Any, *, items: Any, **_: Any) -> IndexCounts: consumed = list(items) @@ -3482,3 +3517,200 @@ def _assert_disk_headroom(*args: Any, **kwargs: Any) -> None: "sweep": 64.0, "other": 9.0, } + + +# --- process-pool extraction (#108) ------------------------------------------- + + +@pytest.mark.unit +def test_default_helper_config_creates_no_extraction_subprocess( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """T9: the shared _config()/_run() helper pair pins extract_processes=1, so + the ~150 run() tests in this file spawn no worker processes and run no + preflight probe. Without this pin every one of them would spawn up to 8 + workers, each re-importing tree_sitter_language_pack -- minutes slower and + flaky under CI load.""" + probe_calls = {"n": 0} + + def _tracking_probe(*args: Any, **kwargs: Any) -> bool: + probe_calls["n"] += 1 + return True + + monkeypatch.setattr("indexer.extract_pool._run_preflight_probe", _tracking_probe) + + code = _run(_config(repos=["acme/widgets"]), _RecordingIndex()) + assert code == 0 + assert probe_calls["n"] == 0 + + +@pytest.mark.unit +def test_pool_engaged_attributes_stream_production_to_parse( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + """T10: phase-timing attribution survives the pool. The pool is built in + run(), OUTSIDE branch_started, so it contributes nothing to any branch's + total by itself; only the per-item production _timed_items wraps is + charged, and it still lands on `parse`, never `other`.""" + clock = _install_fake_clock(monkeypatch) + + class _ClockAdvancingExecutor: + """Advances the fake clock once per submitted batch, proportional to + its file count -- mirrors this file's other `_slow_*` timing fakes, + just at the pool's own submit() seam instead of a direct function + patch.""" + + def submit(self, fn: Any, *args: Any) -> "Future[list[FileExtraction]]": + batch = args[0] + clock.advance(1.5 * len(batch)) + fut: "Future[list[FileExtraction]]" = Future() + fut.set_result(fn(*args)) + return fut + + pool = ExtractionPool(n_processes=0, batch_bytes=10_000_000) # one big batch + pool._n_processes = 2 + pool._executor = _ClockAdvancingExecutor() # type: ignore[assignment] + + idx = _RecordingIndex() + cfg = Settings(semantic_enabled=False) + # Both files must be a recognized language (unlike _DEFAULT_FILES' README.md, + # `lang=None`) so BOTH qualify for the pool and land in the SAME one-file- + # exceeds-nothing batch -- otherwise D5's local short-circuit answers one of + # them without ever reaching the executor, undercounting this assertion. + github = _GitHub( + files={"a.py": b"def f():\n return 1\n", "b.py": b"def g():\n return 2\n"} + ) + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run( + _config(repos=["acme/widgets"]), idx, cfg=cfg, github=github, extraction_pool=pool + ) + assert code == 0 + + phases = _only_phases(caplog) + assert list(phases) == [ + "total", + "resolve", + "download", + "parse", + "embed", + "db", + "sweep", + "other", + ] + assert phases["parse"] == 3.0 # a.py + b.py (the github= override above) x 1.5s, one batch + assert phases["other"] == 0.0 + assert phases["db"] == 0.0 + assert phases["total"] == 3.0 + + +@pytest.mark.unit +def test_extraction_pool_is_shut_down_even_when_fan_out_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Pool lifecycle: built once per run, shut down in the `finally` even when + something inside the fan-out block raises before any branch runs.""" + import indexer.job as job + + shutdown_calls = {"n": 0} + + class _FakePool: + def shutdown(self) -> None: + shutdown_calls["n"] += 1 + + monkeypatch.setattr(job, "build_extraction_pool", lambda config: _FakePool()) + + def _boom_stamps(engine: Any, entries: Any) -> Any: + raise RuntimeError("boom") + + monkeypatch.setattr(job, "_read_stamps", _boom_stamps) + + with pytest.raises(RuntimeError): + _run(_config(repos=["acme/widgets"]), _RecordingIndex()) + + assert shutdown_calls["n"] == 1 + + +@pytest.mark.unit +def test_pool_is_built_before_the_first_db_read(monkeypatch: pytest.MonkeyPatch) -> None: + """R2 (§10.3 of the plan): the preflight probe is the first process spawn in + run(), and it MUST precede any database read or write. A re-entrant + python_wheel_task child under an unguarded __main__ inherits the parent's + sys.argv and would otherwise become a second corpus writer; the probe's + unguarded-__main__ crash (reproduced: RuntimeError from + multiprocessing.spawn._check_not_importing_main) only bounds that risk if + it runs before the child could touch the database. This pins the ordering + directly so a future refactor (e.g. "only build the pool if a branch needs + it") cannot silently move the pool build after _read_stamps without a test + failing -- do not "fix" a failure here by reordering the assertion.""" + import indexer.job as job + + calls: list[str] = [] + + def _tracking_build(config: Any) -> Any: + calls.append("pool") + return ExtractionPool(n_processes=0) + + real_read_stamps = job._read_stamps + + def _tracking_read_stamps(*args: Any, **kwargs: Any) -> Any: + calls.append("stamps") + return real_read_stamps(*args, **kwargs) + + monkeypatch.setattr(job, "build_extraction_pool", _tracking_build) + monkeypatch.setattr(job, "_read_stamps", _tracking_read_stamps) + + code = _run(_config(repos=["acme/widgets"]), _RecordingIndex()) + assert code == 0 + assert calls == ["pool", "stamps"] + + +@pytest.mark.unit +def test_injected_extraction_pool_is_never_shut_down_by_run() -> None: + """Mirrors owns_engine/owns_http: a CALLER-supplied pool is the caller's to + manage, never torn down inside run().""" + shutdown_calls = {"n": 0} + + class _FakePool: + def stream(self, files: Any) -> Any: + return ((pf, FileExtraction(symbols=[], edges=[])) for pf in files) + + def shutdown(self) -> None: + shutdown_calls["n"] += 1 + + code = _run(_config(repos=["acme/widgets"]), _RecordingIndex(), extraction_pool=_FakePool()) + assert code == 0 + assert shutdown_calls["n"] == 0 + + +@pytest.mark.unit +def test_broken_pool_fails_only_the_in_flight_branch_next_branch_recovers( + caplog: pytest.LogCaptureFixture, +) -> None: + """AC3: a shared pool's break fails the branch it broke under; the run + continues and the NEXT branch of the same repo indexes normally -- the + per-branch classification this design maps ExtractionPoolError onto.""" + call_count = {"n": 0} + + class _FlakyPool: + def stream(self, files: Any) -> Any: + call_count["n"] += 1 + if call_count["n"] == 1: + raise ExtractionPoolError("simulated pool break (generation 0)") + return ((pf, FileExtraction(symbols=[], edges=[])) for pf in files) + + def shutdown(self) -> None: + pass + + idx = _RecordingIndex() + github = _GitHub(branches={"acme/widgets": ["main", "feature"]}) + # Default branch ("main") is always first and always included -- see + # indexer.branches.resolve_branches -- so "main" is the branch that breaks + # and "feature" is the one that recovers. + config = _config(repos=["acme/widgets"], branches=["feature"], index_concurrency=1) + + with caplog.at_level(logging.INFO, logger="indexer.job"): + code = _run(config, idx, github=github, extraction_pool=_FlakyPool()) + + assert code == 1 # the run continues but is not clean + assert "failed to index acme/widgets@main" in caplog.text + assert idx.calls == ["acme/widgets"] # only "feature" reached index_repo diff --git a/tests/unit/test_repo_config.py b/tests/unit/test_repo_config.py index e0d0eb4..5861752 100644 --- a/tests/unit/test_repo_config.py +++ b/tests/unit/test_repo_config.py @@ -213,6 +213,41 @@ def test_index_concurrency_out_of_range_raises_config_error(value: int) -> None: assert "index_concurrency" in str(excinfo.value) +# --- extract_processes (#108) ------------------------------------------------- + + +@pytest.mark.unit +def test_extract_processes_defaults_to_none() -> None: + """None means "derive from the runtime" (indexer.extract_pool.derive_process_count) -- + omitting the field is the supported shape, since it predates #108.""" + assert parse_config(_MINIMAL, source="cfg").extract_processes is None + + +@pytest.mark.unit +@pytest.mark.parametrize("value", [1, 4, 8]) +def test_extract_processes_accepts_in_range(value: int) -> None: + raw = b"version: 1\nconnections:\n - type: github\n users: [u]\nextract_processes: %d\n" % ( + value + ) + + assert parse_config(raw, source="cfg").extract_processes == value + + +@pytest.mark.unit +@pytest.mark.parametrize("value", [0, 9]) +def test_extract_processes_out_of_range_raises_config_error(value: int) -> None: + """Matches index_concurrency's `1..8` ceiling -- every parallelism knob in this + repo shares it, and the pool's own sizing derivation already clamps to 8.""" + raw = b"version: 1\nconnections:\n - type: github\n users: [u]\nextract_processes: %d\n" % ( + value + ) + + with pytest.raises(ConfigError) as excinfo: + parse_config(raw, source="cfg") + + assert "extract_processes" in str(excinfo.value) + + @pytest.mark.unit @pytest.mark.parametrize( ("configured", "semantic_enabled", "expected"), From a68d509a16540b3f763005d5b0ec27a297c4d2af Mon Sep 17 00:00:00 2001 From: Tanner <84605639+IceRhymers@users.noreply.github.com> Date: Sun, 26 Jul 2026 13:25:13 -0700 Subject: [PATCH 9/9] indexer: raise semantic worker clamp from 2 to 4 (#109) Refs #109 --- app/config.py | 27 +- config.yaml | 49 ++- docs/perf/issue-109-measurements.md | 483 ++++++++++++++++++++++ docs/runbooks/indexing-parallelism.md | 106 ++++- docs/runbooks/semantic-enablement.md | 55 ++- indexer/ingest.py | 12 + indexer/job.py | 25 +- indexer/repo_config.py | 29 +- scripts/measure_ingest_threads.py | 473 +++++++++++++++++++++ scripts/measure_semantic_memory.py | 574 ++++++++++++++++++++++++++ tests/unit/test_job.py | 126 +++++- tests/unit/test_repo_config.py | 7 +- 12 files changed, 1876 insertions(+), 90 deletions(-) create mode 100644 docs/perf/issue-109-measurements.md create mode 100644 scripts/measure_ingest_threads.py create mode 100644 scripts/measure_semantic_memory.py diff --git a/app/config.py b/app/config.py index 1c5322c..f7a7641 100644 --- a/app/config.py +++ b/app/config.py @@ -93,12 +93,16 @@ class Settings(BaseSettings): # in-process memory rather than a DB lock; exceeding it fails loudly. # # Sized against the ACTUAL buffer cost, not a round number: the vectors are held as - # Python float lists, ~32 B per element (24 B float object + 8 B list pointer), so at - # dim=1024 each chunk costs ~32 KB -- 8000 chunks is ~260 MB of vectors plus ~16 MB of - # chunk text. A larger ceiling (e.g. 50k -> ~1.6 GB) would OOM the job container before - # this loud check could ever fire, which would defeat the point of having a ceiling. - # A repo that legitimately exceeds this needs a temp-table staging path, not a bigger - # buffer. + # Python float lists, ~32 B per element (24 B float object + 8 B list pointer) structural, + # but ~40.1 KB/chunk RESIDENT once measured (issue #109; pymalloc overhead/fragmentation -- + # use this figure for headroom arithmetic) -- 8000 chunks is ~313 MiB of vectors resident + # plus ~16 MB of chunk text. #109 also derived a per-worker chunk-cap ceiling from a full + # container-memory model (~73,300 chunks at the pinned N=2 semantic-worker count this cap + # is evaluated at, ~36,700 at the shipped N=4 -- see docs/perf/issue-109-measurements.md + # §12): a larger ceiling well past that (e.g. 50k, ~1.9 GiB resident) risks OOMing the job + # container before this loud check could ever fire, which would defeat the point of having + # a ceiling. A repo that legitimately exceeds this needs a temp-table staging path, not a + # bigger buffer. # # Scope note (#104): under file-level delta indexing this cap is enforced against # whatever ONE RUN embeds (changed/new + membership-only files), not a branch's whole @@ -112,11 +116,12 @@ class Settings(BaseSettings): # which bounds file ingestion, not embedding-chunk granularity. semantic_chunk_max_tokens: int = 512 - # In-flight embedding requests per worker (#107). The indexer clamps to 2 workers when - # semantic is on (indexer/repo_config.py:effective_workers), so total in-flight gateway - # requests are workers x concurrency: 2 x 4 = 8 at this default, 2 x 8 = 16 at the - # config.yaml-enforced ceiling of 8 -- both under the SDK's 20-connection pool - # (pool_block=True, so exceeding it would silently serialize rather than error). Setting + # In-flight embedding requests per worker (#107). The indexer clamps to 4 workers when + # semantic is on (indexer/repo_config.py:effective_workers -- issue #109 raised this from + # 2), so total in-flight gateway requests are workers x concurrency: 4 x 4 = 16 at this + # default, 4 x 8 = 32 at the config.yaml-enforced ceiling of 8 -- the latter now EXCEEDS + # the SDK's 20-connection pool (pool_block=True, so exceeding it silently serializes + # rather than erroring, so this is a real-concurrency cap, not a correctness one). Setting # this to 1 restores today's fully serial embed() and spawns no thread pool. semantic_embedding_concurrency: int = 4 diff --git a/config.yaml b/config.yaml index 28eff07..36fd6e5 100644 --- a/config.yaml +++ b/config.yaml @@ -12,8 +12,13 @@ version: 1 # threads — the tree walk is GIL-serialized), which is why extraction now runs # in its own shared process pool instead — see extract_processes below. Raising # this knob buys disk-bound repo fan-out, not extraction throughput. -# With semantic indexing enabled this is clamped to 2 — a MEMORY bound, since -# embedding materialises a whole repo's chunks (~0.5-0.8 GB per worker). +# With semantic indexing enabled this is clamped to 4 (issue #109 raised it from +# 2, after re-deriving the memory model and confirming empirically against the +# live dev job at N=4: peak self+children RSS landed at ~83% of the 0.7*container +# memory budget, comfortably under and with more margin than N=2's own ~90%) — a +# MEMORY bound, since embedding materialises a whole repo's chunks (structural +# ~32 KB/chunk, resident ~40.1 KB/chunk measured; see effective_workers' +# docstring and docs/perf/issue-109-measurements.md for the full derivation). # index_concurrency: 4 # How many worker PROCESSES the job uses to extract symbols/edges (issue #108). @@ -75,13 +80,23 @@ connections: # effective cap as `per-repo override OR global`. # # Mind the memory math before raising one: buffered vectors are ~32 KB/chunk -# (dim=1024, Python float-list storage), so 8000 ≈ 260 MB resident for the -# duration of that repo's write. With semantic on, at most 2 workers run -# concurrently (indexer/repo_config.py's effective_workers clamp), so a large -# override multiplies straight into the job container's peak memory — e.g. two -# repos overridden to 20000 concurrently is ≈1 GB just in vectors, on top of -# the base per-worker cost. A repo that legitimately needs far more than that -# needs the temp-table staging path (follow-up), not a bigger override. +# structural (dim=1024, Python float-list storage) but ~40.1 KB/chunk RESIDENT +# (measured, issue #109 — includes pymalloc overhead/fragmentation; use this +# figure for headroom arithmetic), so 8000 ≈ 313 MiB resident for the duration of +# that repo's write. With semantic on, at most 4 workers run concurrently +# (indexer/repo_config.py's effective_workers clamp, raised from 2 by #109), so +# a large override multiplies straight into the job container's peak memory — +# e.g. two repos overridden to 20000 concurrently is ≈1.6 GB just in vectors, on +# top of the base per-worker cost. A repo that legitimately needs far more than +# that needs the temp-table staging path (follow-up), not a bigger override. +# +# The derived per-worker chunk-cap ceiling, from a full container-memory model +# (issue #109; docs/perf/issue-109-measurements.md §12 — pinned at N=2 there +# only to break a circularity in solving for C from a formula whose dominant +# term IS C, not a claim about the adopted concurrency): ≈73,300 chunks at +# N=2, ≈36,700 at the shipped N=4 (both halve/double with N). The current +# global default of 8000 uses well under a quarter of either budget, so it is +# NOT the binding constraint and was left unchanged. # semantic_max_chunks_per_repo: # "acme/huge-monorepo": 20000 @@ -98,20 +113,24 @@ connections: # (the default is 8000). It is NOT the per-repo `semantic_max_chunks_per_repo` MAP # above — that spot-overrides individual repos and still wins over this global. Use # this to raise the floor everyone inherits; use the map for the outliers. The same -# ~32 KB/chunk memory math and 2-worker clamp above apply here, magnified: raising -# the global lifts the buffer cost for EVERY concurrently-indexing repo at once. +# ~32 KB structural / ~40.1 KB resident per-chunk memory math and 4-worker clamp +# above apply here, magnified: raising the global lifts the buffer cost for EVERY +# concurrently-indexing repo at once. # # `enabled: false` makes the job a true semantic no-op (no embedder built, no -# chunking, the 2-worker memory clamp not applied) even if the env says enabled — +# chunking, the 4-worker memory clamp not applied) even if the env says enabled — # the fastest way to turn semantic off for the job alone. # # `embedding_concurrency` (#107) is in-flight embedding requests PER WORKER, sent # via a ThreadPoolExecutor that preserves submission order — vectors always come # back in the order their texts were sent, regardless of which request finishes # first. Total in-flight gateway requests for the job is workers x concurrency: -# 2 x 4 = 8 at this default, 2 x 8 = 16 at the max of 8, both under the SDK's -# 20-connection pool. Set to 1 to restore fully serial embedding (no thread pool -# spawned at all) if you need to roll back. +# 4 x 4 = 16 at this default, 4 x 8 = 32 at the max of 8 — the latter now EXCEEDS +# the SDK's 20-connection pool (issue #109 raised workers from 2 to 4; this +# combination was not possible before). Lower embedding_concurrency if raising it +# alongside a near-ceiling index_concurrency. Set embedding_concurrency to 1 to +# restore fully serial embedding (no thread pool spawned at all) if you need to +# roll back. # semantic: # enabled: true # max_chunks_per_repo: 8000 diff --git a/docs/perf/issue-109-measurements.md b/docs/perf/issue-109-measurements.md new file mode 100644 index 0000000..edc2caf --- /dev/null +++ b/docs/perf/issue-109-measurements.md @@ -0,0 +1,483 @@ +# Issue #109 — re-derive worker, disk, and memory limits: measurements + +**Unit convention, stated up front (mixing anchors is a real risk in this doc):** +every RSS/vector/memory-model figure below is **binary** (KiB/MiB = 1024-based), +matching `resource.ru_maxrss` (Linux: KB = 1024 bytes) and this repo's existing +`docs/perf/issue-108-measurements.md` convention. The one exception is +`MAX_EXTRACTED_BYTES = 2_000_000_000`, which is **decimal** (2 GB = 2,000,000,000 +bytes) by the source constant's own definition — any comparison against it is +converted explicitly, never left implicit. + +--- + +## 1. Environment + +**Local box** (measurements in §2–§4): Linux, 12 cores, ~15.5 GB RAM, Python +3.12.13 (uv-managed `.venv`; the shell's own default is 3.14). `/tmp` is tmpfs +(7.8 GB) — disk-backed measurements used `/` (nvme0n1p2, 318 GB free) instead. + +**Databricks dev serverless container** (§5, `M` and W4): read directly from +inside the deployed job via a temporary probe (job schedule stayed **PAUSED** +throughout; only manually-triggered one-time runs executed). Two separate +container instances were observed across two probe attempts, with a **>2x +spread** in reported memory — see §5.1. + +--- + +## 2. AC2 — disk (E1): already correct on the base, verified not re-derived + +Per the plan's §0.1, #106 already landed the disk half of #109. Verified by +citation, not rewritten: + +- `indexer/fetch.py`: `REQUIRED_FREE_BYTES == MAX_TARBALL_BYTES == 500_000_000`. +- The guard message carries the real number (`need 500000000 ...`). +- `docs/runbooks/indexing-parallelism.md` §3 and `config.yaml` already read 0.5 + GB/worker (1→0.5, 2→1, 4→2, 8→4 GB). +- The tarball is the ONLY on-disk artifact (`indexer.ingest.iter_tar_source_files` + streams in memory, never extracts) — confirmed by direct code reading, not a + fresh sampling run (§0.1 forbids "restating a correct 8-worker figure" as new + work). +- W3 (observed on the live dev job, both arms): `local disk at /tmp: 64.0 GB + free of 89.1 GB total` — **disk is not a binding constraint at any allowed + `index_concurrency` (up to 8, 4 GB peak)**. + +**AC2: satisfied, unchanged.** + +--- + +## 3. M1 — the memory model's coefficients (E3(b–f)) + +`scripts/measure_semantic_memory.py`, run against 4 corpora (this repo, +`flask`, `requests`, `django`), driving the REAL semantic path +(`iter_tar_source_files` → delta narrowing → `_precompute_chunk_writer`) with a +stub embedder returning **distinct** floats per chunk (per §2.2's trap — a +shared-cached-float stub understates resident vector cost ~4x). + +| Corpus | alpha | gamma | d (bytes/chunk) | +|---|---|---|---| +| databricks-code-search | 2.1142 | 0.5571 | 2711.6 | +| flask | 2.044 | 0.4693 | 1644.2 | +| requests | 3.382 | 0.0 (see below) | 3747.7 | +| django | 1.5911 | 1.2627 | 1775.5 | + +`requests`' gamma measured as 0 — a chunk delta small enough (chunk_count=413) +to fall below this box's RSS sampling granularity (allocator/page-granularity +noise), not evidence chunking is free for that corpus. + +alpha: avg=2.2828, **max=3.382** (n=4). gamma: avg=0.5723, **max=1.2627** (n=4). +**Decision: use MAX-observed coefficients** (larger → smaller/safer derived +limits, larger/stricter `P_worst`) as the primary, conservative input; average +reported alongside for context (§8 judgement call, self-consistently applied +everywhere it's used). **`alpha + gamma = 4.6447`** (max of each, not the max +corpus's sum — the model treats them as independently-conservative). + +`V_cap` (resident vector cost, 8000 × 1024 distinct-float vectors): **40.102 +KB/chunk** (this session's re-measurement; close to planning's 40.8 KB/chunk — +both measure the same thing, structural is 32.0 KB/chunk exact +(`1024 × (8B pointer + 24B PyFloat)`), not "corrected" by this re-measurement, +per §2.2's own instruction not to). + +`R_proc` = 121 MB/process (#108's own isolated measurement — reused, not +re-run, per the plan's L4). + +`P_fixed`: a bare-interpreter floor probe gave 38,372 KB → +`import indexer.job` +(pulls SQLAlchemy/databricks-sdk) → 63,048 KB → +a throwaway SQLAlchemy engine → +67,628 KB. This is a **floor** (~68 MB) — it excludes the real pool_size-scaled +connection pool and Databricks SDK client state a live job carries. **Adopted +P_fixed = 300 MB** (the plan's own conservative worked-example value), noting +the ~68 MB floor as a cross-check, not a replacement. + +### 3.1 A methodology finding not anticipated by the plan: fork-time COW RSS contamination of `RUSAGE_CHILDREN` + +M1's stage-3 (N-concurrency) sweep reported `after_children_kb` **numerically +identical** to `after_self_kb` at every N (e.g. N=1: self=1,046,984, +children=1,046,984). Root-caused directly (scratch scripts, not committed): +when the extraction pool's `spawn` workers are forked/exec'd **after** a +repo-worker's `_precompute_chunk_writer` has already ballooned the calling +thread's RSS to ~1 GB — exactly production's real call order in +`indexer/job.py` — each child's `ru_maxrss` (read later via `wait()` / +`RUSAGE_CHILDREN`) captures the **fork-time COW snapshot of the parent's +then-current RSS**, not the child's real post-exec working set. Verified +directly: draining the pool BEFORE ballooning gives children ~144–150 MB +(matching #108's own R_proc ~121 MB order of magnitude); draining AFTER gives +children ≈ self's contemporaneous value — a ~7–9x inflation with **no +corresponding real memory pressure** (COW pages are shared, billed once by the +cgroup, not per-process). + +**Implication:** this affects M1's own stage-3 "children" column, and by the +same mechanism, the production `peak rss: self=... children=...` instrumentation +on Arms A/B (`indexer/job.py`'s new n5 log line) whenever the pool is +(re)spawned after chunk-writer inflation. It is a `ru_maxrss` **measurement +artifact**, not evidence of doubled real memory. The `P_worst` model itself is +unaffected in its primary form because `R_proc` is sourced from #108's own +isolated measurement, not from this contaminated figure — but Arm A/B's +observed `children=` numbers below should be read as **upper bounds, not +literal per-process costs**. + +Stage-3 self-deltas (uncontaminated — `self` reflects only the process's own +allocations) at N=1..4, worst-case corpus (django, first-index/gate-closed): +N=1: 983,036 KB; N=2: 1,828,428 KB; N=3: 2,582,480 KB; N=4: 3,216,572 KB — +**sub-linear**, consistent with page-cache/tarball-decompression sharing across +threads, not a red flag. + +**Caveat on these specific numbers (found and fixed post-hoc in review, not +re-measured):** at measurement time, `scripts/measure_semantic_memory.py`'s +stage-3 worker discarded `_precompute_chunk_writer`'s return value before +calling `pool.stream()`, so each thread's vectors were collectable before (or +concurrently with) sibling threads' peaks — understating true N-way +concurrent residency relative to production, which holds `chunk_writer` alive +across the whole write window. The script is fixed in this PR (the return +value is now held alive across `pool.stream()` and explicitly `del`eted after, +matching stage 1's pattern) for future use, but the N=1..4 numbers above +**were not re-measured** against the fix, since **no decision in this PR rests +on them** — the adopted N=4 decision comes from the `P_worst` model (§6) and +the real Arm A/B live-job runs (§8–§10), not from this local sub-measurement. +Treat the sub-linearity finding above as directional, not load-bearing. + +--- + +## 4. B — the three anchors + +- **`B_prod`** (real corpus, unnesting `files.branches`): top row `repo_id=46` + (`IceRhymers/opencode`, branch `dev`), `src_bytes = 31,778,187` (~30.3 MiB). +- **`B_obs`** (measurement corpus, re-measured directly): `opencode@dev` = + 33,015,231 bytes decoded source (18,617 chunks, 4,759 files) — slightly + higher than `B_prod` (encoding/whitespace differences between the two + measurement paths). +- **`B* = max(B_prod, B_obs) = 33,015,231 bytes ≈ 31.49 MiB`** — the anchor used + for every decision below. +- **`B_ceil = MAX_EXTRACTED_BYTES = 2 GB`** — theoretical, loose, **never** used + to drive a decision (only to illustrate why a naive `B_ceil`-anchored model + would falsely condemn the status quo — see below). + +--- + +## 5. `M` — the container memory ceiling (E3(a)) + +Two container instances observed across two probe attempts on the same job: + +| Attempt | task_run_id | `cgroup_v1_memory.limit_in_bytes` | `sched_getaffinity` | Outcome | +|---|---|---|---|---| +| 1 | 523922694112794 | 8,385,462,272 (~7996 MiB) | 4 | **OOM-killed** during the allocate-bracket, last logged step `cumulative_mb=12800` | +| 2 | 89875228623550 | 24,706,547,712 (~23.0 GiB), `MemTotal` 32,264,556 kB | 4 | Ran its bracket to a designed 16,384 MB cap without dying (never exercised further) | + +**`cgroup_v2_memory.max`/`memory.high` both unreadable** (`FileNotFoundError`) +on this runtime — the repo's own `extract_pool.py::_cgroup_cpu_quota()` v2-only +assumption does not hold for `memory.max`; the v1 fallback +(`/sys/fs/cgroup/memory/memory.limit_in_bytes`) was required. + +**Attempt 1's cgroup read was demonstrably a MISREAD**: the container died at +`cumulative_mb=12800` (the bracket's last logged step before the kill), i.e. +the real ceiling sits in `(12800, ~13056] MiB` — **~60% higher** than the +cgroup-reported ~7996 MiB. + +**`M = 12800 MiB`** (the smaller, real, empirically-grounded ceiling from +attempt 1 — chosen conservatively over attempt 2's larger, undead container, +per the resumption brief's instruction to use the smaller real ceiling for +safety). **`0.7 × M = 8960 MiB`** — the budget used throughout. + +**`W4` (container CPU count) = 4** (`len(os.sched_getaffinity(0))`, both +attempts agree) — sourced from this probe, per the plan's design, breaking the +apparent circularity between E4 (needs W4) and E5 (Arm B, which would +otherwise be W4's only source). + +### 5.1 First-order finding: container sizing is unstable across attempts + +A **≥2x spread** in effective memory ceiling was observed between two +instances of the *same* job on *unspecified* dev serverless compute (~8 GiB vs. +~23.0 GiB reported; real ceilings both plausibly larger than reported). This is +reported as a finding, not resolved — the smaller, conservative number is what +every downstream decision uses. + +--- + +## 6. `P_worst` — the model, evaluated at the standard cap + +``` +P_worst(N) = N × max( (alpha+gamma)·B_breach, + (alpha+gamma)·(d×C) + V_cap·C ) + + extract_processes × R_proc + P_fixed +``` + +At the **standard/default global chunk cap `C = 8000`** (ordinary unmodified +production, NOT the measurement corpus's per-repo overrides — see §7 for why +those diverge), with the MAX coefficients above (`alpha+gamma = 4.6447`, +`d = 3747.7`, `V_cap = 40.102 KB/chunk`, `R_proc = 121 MB`, +`extract_processes = 4` — confirmed live in the priming log's "symbol +extraction: 4 process(es)"), `P_fixed = 300 MB`: + +- Breach term at `B* = 31.49 MiB`: `(a+g)·B* ≈ 146.3 MiB`/worker. +- Under-cap term at `C = 8000` (the regime that wins under the standard cap — + vectors dominate the uncapped terms ~8x at the cap, per the plan's §2.3): + `(a+g)·d·C ≈ 132.8 MiB` + `V_cap·C ≈ 313.3 MiB` = **446.1 MiB**/worker. +- `max(146.3, 446.1) = 446.1 MiB`/worker. + +`P_worst(3, B*) = 3×446.1 + 4×121 + 300 = 2122.3 MB` +`P_worst(4, B*) = 4×446.1 + 4×121 + 300 = 2568.4 MB` + +Both `<< 0.7M = 8960 MiB` (margin ~76% at N=3, ~71% at N=4). Cross-checked with +AVERAGE coefficients (`alpha+gamma = 2.855`, avg `d = 2469.75`): +`P_worst(4) = 2252.4 MB` — same conclusion; **not sensitive to the max-vs-avg +coefficient choice.** + +**Illustrative-only counter-example (never used to decide anything): at +`B_ceil = 2 GB` (= 1907.3 MiB binary)** the breach term is +`(alpha+gamma)·B_ceil = 4.6447 × 1907.3 ≈ 8859.1 MiB`/worker — the breach +regime wins by **≈19.9x** over the under-cap term (446.1 MiB, §6), giving +`P_worst(2, B_ceil) = 2×8859.1 + 484 + 300 ≈ 18,502 MiB` — i.e. the model +would falsely condemn TODAY's clamp=2 status quo many times over if evaluated +at the loose theoretical bound instead of the real observed `B*`. This is +exactly why §3.3(c) of the plan anchors the decision on `B*`, never `B_ceil`. +(The planning-stage worked example in the plan document itself used its own +pre-`M1` coefficients, ~2.35 rather than the measured 4.6447, and got a +smaller — still condemning — ~10.2 GiB; both versions support the same +qualitative point, so both numbers are recorded here for provenance: +whichever coefficient set is used, `B_ceil` is not a fit substitute for `B*`.) + +**Decision (model-only, pre-Arm-A): branch (i), RAISE — both N=3 and N=4 clear +`0.7M` with large margin. Adopt N=4 (the largest passing N), contingent on Arm +B completing cleanly** — modelling alone is not sufficient per the plan's AND +condition. + +--- + +## 7. Re-evaluation against Arm A's own observed peak (§3.3(c).iii) + +The standard-cap model above does **not** reflect what Arm A actually ran, +because the measurement corpus's per-repo chunk-cap overrides (opencode = +22,340, ≈2.8x the standard 8000 cap — see §9) make this specific corpus far +more memory-expensive than "standard production": Arm A's real self+children +was ~4.75x higher than the standard-cap model's own N=2 prediction (1676 MiB), +because the override intentionally lets opencode use ~874.9 MiB of vectors +instead of breaching into the cheap ~146 MiB breach regime. Expected (flagged +in the plan's §3.2 as a likely consequence), not a bug — but it means only +Arm A's own empirical number, not the standard-cap model, can gauge Arm B's +real risk on this corpus. + +--- + +## 8. Arm A — before (clamp=2, TODAY's shipped code) + +Both runs: 22 repos resolved (19 retained + 3 explicit), `symbol extraction: 4 +process(es) (spawn); pool preflight ok`, `local disk ...; 2 worker(s) x 0.5 GB +peak`, `semantic enabled: clamping index_concurrency 4 -> 2`, 22/22 branches ok, +0 skipped/conflicts/failed, 0 repos purged, no `degraded semantic coverage` +WARNING, no 429/retry lines. + +| | Run 1 (`489590218152767`) | Run 2 (`19039699744754`) | +|---|---|---| +| wall (execution_duration) | 350.5s | 342.4s | +| `peak rss: self=... children=...` | self=6,930,768 KB children=1,247,176 KB | self=7,059,172 KB children=1,252,284 KB | +| self+children | 8,177,944 KB ≈ **7986.3 MiB** | 8,311,456 KB ≈ **8116.7 MiB** | +| % of `0.7M` (8960 MiB) budget | 89.1% | 90.6% | + +The two runs agree within 1.6 percentage points — **not a fluke: at TODAY's +clamp=2, this corpus already consumes ~89–91% of the safety budget.** Best-of +(faster wall): run 2, 342.4s. Worse-of (higher memory, used as the conservative +anchor going forward): run 2, 8116.7 MiB. + +**§3.3(c).iii re-evaluation:** this real number is ~4.75x the standard-cap +model's prediction (§7) — the standard-cap model cannot gauge Arm B's risk on +this corpus. Direct reasoning from Arm A's own peak: this 22-repo corpus has +exactly 3 memory-heavy repos (opencode ~874.9 MiB vectors, claw-code ~147.9 MiB, +nanoclaw ~30 MiB), and overlap is capped by there being only 3 of them +regardless of N — so N=3/N=4 were estimated to land close to Arm A's own +number (perhaps +50–150 MiB), i.e. plausibly under budget but with a **thinner +margin (~85–90% utilized)** than the standard-cap model implied. **Decision: +still target N=4** for the single Arm B attempt (the largest of {3,4}, and +nothing in the refined reasoning favors N=3 specifically), with full awareness +of the thinner real margin and a non-negligible chance of failure — a valid, +reportable outcome either way per §6, not a trigger to retry at a different N. + +--- + +## 9. Arm B — after (clamp=4, the derived change) + +Applied the one-line change (`indexer/repo_config.py::effective_workers`: +`min(index_concurrency, 2)` → `min(index_concurrency, 4)`), redeployed +(`databricks bundle deploy -t dev`), cleared stamps, warmed Lakebase (>300s +`SELECT 1` loop immediately before each timed run — see §10 for a warm-up +reliability note), ran twice. + +Both runs: 22/22 branches ok, 0 skipped/conflicts/failed, 0 repos purged, +**no** `semantic enabled: clamping ...` line (correct: `index_concurrency=4` +now equals the clamp, so `effective_workers` is a no-op passthrough — matching +§11's conclusion that the `index_concurrency` default itself did not need to +move), `local disk ...; 4 worker(s) x 0.5 GB peak`, `symbol extraction: 4 +process(es) (spawn); pool preflight ok`, **zero** WARNING/ERROR/retry/429 lines +anywhere in either 191-line log. + +| | Run 1 (`752462522914821`) | Run 2 (`785234477657138`) | +|---|---|---| +| wall (execution_duration) | 307.8s | 305.5s | +| `peak rss: self=... children=...` | self=6,128,880 KB children=1,445,408 KB | self=6,225,624 KB children=1,429,280 KB | +| self+children | 7,574,288 KB ≈ **7396.8 MiB** | 7,654,904 KB ≈ **7475.5 MiB** | +| % of `0.7M` (8960 MiB) budget | 82.6% | 83.4% | + +Both runs agree within 1 percentage point — not a fluke. Best-of (faster +wall): run 2, 305.5s. + +**Surprising but real: Arm B's total peak is LOWER than both Arm A runs, and +its margin (~17%) is MORE comfortable than Arm A's own (~9–11%).** `ru_maxrss` +is a same-process high-water mark, so this is a genuine peak-memory +observation, not a modelling artifact — Arm B's own `self` component (6.13–6.23 +million KB) is genuinely lower than either Arm A `self` value (6.93–7.06 +million KB). Consistent with (not contradicting) §8's reasoning: this corpus +has only 3 memory-heavy repos, and with 4 workers instead of 2 their embedding +windows are *less* likely to bunch up at the tail (more workers drain the +22-repo queue faster and spread the 3 heavy repos across more concurrent slots +with shorter individual overlap) — a property specific to this corpus's +repo-size distribution, **not** a general "N=4 always costs less than N=2" +claim. + +**No stop condition (§6 of the plan) was triggered anywhere in the Arm B +sequence.** + +## 10. AC1 — before/after table + +| | Arm A (clamp=2, best-of wall) | Arm B (clamp=4, best-of wall) | +|---|---|---| +| wall | 342.4s | 305.5s (**−10.8%**) | +| peak self+children | 8116.7 MiB (90.6% of budget) | 7475.5 MiB (83.4% of budget) | + +**FINAL DECISION (AC1 + AC3): adopt N=4.** Both the `P_worst` model (§6, before +any arm ran) and two clean, mutually-consistent empirical Arm B runs (§9) agree. + +**Lakebase CU state**: `resources/lakebase.yml` pins +`autoscaling_limit_min_cu: 0.5`, `autoscaling_limit_max_cu: 4`, +`suspend_timeout_duration: 300s`. Warmed with a `SELECT 1` loop for the full +300s+ immediately before every timed run (priming, Arm A ×2, Arm B ×2 — 5 +warm-ups total), confirmed complete each time via elapsed wall-clock (302.3s / +301.6s measured, not assumed). + +**Budget/429 posture**: 5 full semantic-index runs total (1 priming + 2 Arm A + +2 Arm B) against the real, paid AI Gateway embedding endpoint. **Zero** 429s or +`databricks.sdk.retries` entries observed in any run's log; zero `degraded +semantic coverage` WARNINGs. + +--- + +## 11. `index_concurrency` default (E4/M2 — the ingest-thread-scaling question) + +`scripts/measure_ingest_threads.py`, repos: fastapi, sqlalchemy, sympy, django. + +- N=2 idle: 1.26x (ambiguous band, ≥1.6 parallelizes / ≤1.1 doesn't). +- N=4 idle: **1.20x** (NO-SCALE — ≤1.2 threshold, exactly at the boundary). +- N=2 pool-live: 1.20x (ambiguous). +- N=4 pool-live: **1.15x** (NO-SCALE). +- Component breakdown confirms the GIL-bound prior (§0.4 of the plan): + `tf_next`/`fh_read` cumulative time grows super-linearly with N (contention); + `decode` stays roughly flat — consistent with a mostly GIL-held pass. + +**Per the plan's §3.3(d) rule 2** ("the default may rise to the smallest value +≥ N that M2 shows a gain for" once the clamp itself rises): M2 shows **no** +real gain at 4 threads. **`index_concurrency`'s default stays 4** — and no +change was even needed: it was already 4 (`config.yaml`'s commented-out +default), so raising the semantic clamp from 2 to 4 alone makes +`effective_workers = min(4, 4) = 4` with zero separate change to +`index_concurrency` itself. + +--- + +## 12. The two derived byte limits (§3.3(f)) — documented, neither changed + +Both solved at the **pinned N=2** operating point (methodological, to break the +circularity of solving for `B` from a model whose dominant term is `B` — not a +claim that N=2 is the adopted concurrency, which is 4 per §10). +`RHS = 0.7M − extract_processes·R_proc − P_fixed = 8960 − 484 − 300 = 8176 MiB`. + +**(f1) `MAX_EXTRACTED_BYTES` (breach regime, V=0):** +`B ≤ RHS / (N·(alpha+gamma)) = 8176 / (2×4.6447) ≈ 880.1 MiB ≈ 922.9 MB +(decimal)`. Current value: `2_000_000_000` bytes = 1907.3 MiB (decimal 2 GB). + +**(f2) The chunk cap `C` (under-cap regime):** +`C ≤ RHS / (N·((alpha+gamma)·d + V_cap))`. Denominator = +`4.6447×3747.7 + 41,064.4 B ≈ 58,471 B ≈ 57.1 KiB`/unit-of-cap. `C ≤ +8176×1024×1024 / (2×58,471) ≈ 73,311 chunks`. Current global default: 8000 +(**~11% of the derived bound** — nowhere near binding). + +**Neither value was changed.** `B* = 31.49 MiB` is only **3.6%** of the derived +`f1` bound (880.1 MiB) — no branch in the measurement corpus approaches either +the current 2 GB constant or the derived ~880 MiB one, so this run gives **no +empirical signal** to justify lowering a constant whose breach fails the branch +and closes the whole run's reconciliation checkpoint (`job.py`'s +`_decide_reconciliation` requires `failures == 0`). Per the plan's explicit +fallback ("otherwise document the derived value and keep 2 GB with a +comment"), both derived values are recorded here and in the source comments +(`indexer/ingest.py`, `config.yaml`) for the next time this needs re-deriving, +but the live constants are unchanged. The chunk cap `C=8000` also stays — +it uses only ~11% of its own derived budget, so it was never a candidate for +change either. + +**Which regime binds at `B*`:** the **breach** regime, at the standard cap. +`B*` (~31.49 MiB, opencode@dev's real source size) exceeds `d×C` at the +standard `C=8000` under either the conservative modelling `d` (3747.7 B/chunk +→ threshold ≈28.6 MiB) or opencode's own real chunking density (≈1773 +B/chunk → threshold ≈13.5 MiB) — i.e. **a repo the size of `B*` would actually +breach the standard 8000-chunk cap**, which is exactly why the measurement +corpus needed opencode's per-repo override (22,340, §9's context) to avoid +degrading it during Arm A/B. This makes `f1` (the breach-regime derived bound, +not `f2`) the relevant limit to check `B*` against — already done above: +`B*` is only 3.6% of `f1`'s ~880 MiB, comfortably clear. `f2`'s under-cap +regime governs a *different* class of repo: one that legitimately reaches the +cap without a bigger override (§6's worked "vectors dominate ~8x" example), +which no repo in this measurement corpus represents at the standard cap. + +--- + +## 13. E7 — the connection-pool property, confirmed and pinned + +`pool_size == effective_workers` (raised to 4), `max_overflow=0`. Holds by +**sequencing, not construction**: `indexer/job.py`'s `with engine.connect() as +shas_conn:` (the advisory `shas_fn` read, #104) opens and closes a second, +short-lived connection strictly BEFORE `_precompute_chunk_writer`/embedding — +confirmed by direct code reading (job.py:1298 area) and pinned by a new test, +`tests/unit/test_job.py::test_shas_fn_connection_closes_before_embedding_starts`, +which counts currently-open connections and asserts exactly 1 while `shas_fn` +runs and 0 by the time embedding starts. No live-job evidence of Lakebase +connection pressure at N=4 (no `QueuePool` timeouts in either Arm B run), but +the ceiling itself did double (2→4 concurrent connections on the semantic +path) — worth knowing if a future Lakebase compute-size change is on the table +alongside a semantic-path concurrency change. + +Also found (documented, not a blocker): `embedding_concurrency` at its +`le=8` config ceiling combined with the new clamp of 4 now yields `4×8=32` +in-flight gateway requests — **exceeding** the SDK's 20-connection pool for +the first time (`2×8=16` stayed under it before). `pool_block=True` means this +degrades to silent serialization, not an error, so it is not a correctness +issue, but is flagged in `config.yaml`, `app/config.py`, and both runbooks. + +--- + +## 14. The semantics tripwire (§5.1 of the plan) + +`tests/unit/test_semantics_version_tripwire.py::test_semantics_change_bumps_the_index_semantics_version` +**fails on this PR**, as expected and pre-documented by the plan: `SEMANTICS_PATHS` +includes `indexer/ingest.py`, which this PR touches (the `MAX_EXTRACTED_BYTES` +comment, §12) and which — per the test's own module docstring — is *added* +relative to `origin/master` locally (postdates master, landed by #106, +unmerged past this integration branch), producing an expected false positive +this test's own docstring names verbatim. **Per the plan's §5.1, this is +explicitly NOT a stop condition.** `INDEX_SEMANTICS_VERSION` was **not** +bumped and `indexer/ingest.py` was **not** removed from `SEMANTICS_PATHS` — +this section is that required "say so in the PR." + +--- + +## 15. Limitations + +- **Prod** is out of scope and unreachable; dev serverless is the labelled + proxy throughout. +- **Corpus scale**: the real dev corpus (19 repos, 1142 files) is far too + small on its own; the measurement corpus (+opencode/nanoclaw/claw-code) is a + labelled proxy, not production traffic. +- **`M`'s instability** (§5.1): a ≥2x spread was observed between two container + instances of the same job. The smaller, conservative reading was used + throughout; the real production ceiling could be substantially larger. +- **`RUSAGE_CHILDREN` COW contamination** (§3.1): Arm A/B's `children=` figures + are upper bounds, not literal per-process costs, whenever the pool is + (re)spawned after chunk-writer inflation (the real, unavoidable production + call order). +- **`B*` is a sample max** over an executor-chosen corpus (§8.3 of the plan): + passing at N=4 on this corpus is necessary, not sufficient evidence for + every possible future repo. diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index ee82ebd..fc7cbd0 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/docs/runbooks/indexing-parallelism.md @@ -324,9 +324,20 @@ Only the first of the two byte caps is a **disk** cap. Since #106 the tarball is streamed once, in memory, and is never extracted, so `MAX_EXTRACTED_BYTES` is a **work** cap — a decompression-bomb guard on how much content one branch may pull out of its archive — and it lives in `indexer/ingest.py`, beside its only -consumer, rather than in `indexer/fetch.py`. The two therefore no longer sum: -the compressed tarball is the only artifact on disk, so peak local disk is -`index_concurrency` × 500 MB: +consumer, rather than in `indexer/fetch.py`. #109 re-derived both this and +`semantic_max_chunks_per_repo`'s global default as **memory** limits (pinned at +the N=2 operating point, to avoid the circularity of solving for `B` from a +model whose dominant term is `B`): a derived breach-regime ceiling of ~880 MiB +(vs. the current 2 GB) and a derived chunk-cap ceiling of ~73,300 chunks (vs. +the current 8000). **Neither was changed**: no repo in the #109 measurement +corpus approached either the current or the derived byte ceiling (largest +observed branch: ~31.5 MiB, ~3.6% of the derived bound), so this run gives no +empirical signal either way, and lowering `MAX_EXTRACTED_BYTES` closes the +whole run's reconciliation checkpoint on breach — too large a blast radius to +change on an untested corpus. Both derived values are recorded here and in +`docs/perf/issue-109-measurements.md` for the next time this needs re-deriving. +The two therefore no longer sum: the compressed tarball is the only artifact on +disk, so peak local disk is `index_concurrency` × 500 MB: | `index_concurrency` | Peak local disk | |---|---| @@ -345,27 +356,68 @@ Raise `index_concurrency` only for repo-level (disk-bound) fan-out; raise `extract_processes` for CPU-bound extraction throughput. The 4 GB disk figure above is still a hard, linear, unavoidable cost of `index_concurrency` alone. (#106 lowered these numbers by 5x but deliberately did **not** move the default -of 4; re-deriving it is #109's job.) - -**Semantic indexing clamps the pool to 2**, regardless of `index_concurrency`. -That clamp is a *memory* bound, not a CPU one: embedding materialises a whole -repo's chunks in memory (~0.5-0.8 GB per worker). The clamp is logged: +of 4; #109 re-derived it and left it at 4 -- see next.) + +**The `index_concurrency` default itself stays 4 (#109 measured, did not just +inherit, this).** The per-thread ingest pass added by #106 +(`iter_tar_source_files`) mixes GIL-bound Python (the `tf.next()` walk, +`fh.read()`, the NUL-strip, `str.decode` -- which does **not** release the GIL) +with one GIL-releasing step (`zlib.decompress`). Measured thread-scaling on this +pass: **1.15-1.20x at 4 concurrent threads**, both idle and with the extraction +pool live -- short of the 2.0x threshold this repo's measurements use to call +something "parallelizes", and consistent with a mostly GIL-held pass rather +than a genuinely parallel one. No default change is warranted on CPU grounds; +the semantic clamp raise to 4 (above) happens to make `effective_workers` equal +`index_concurrency` at the default with zero separate change needed. + +**Semantic indexing clamps the pool to 4**, regardless of `index_concurrency` +(issue #109 raised this from 2). That clamp is a *memory* bound, not a CPU one: +embedding materialises a whole repo's chunks in memory (~32 KB/chunk structural, +~40.1 KB/chunk resident, measured). **The clamp gates `index_concurrency` +itself: raising the configured default above 4 is a no-op on the semantic path +until the clamp moves too** (`effective_workers` is `min(index_concurrency, +clamp)`) — this is why #109 could not treat the two knobs independently. The +clamp is logged only when it actually reduces the configured value: ``` -INFO indexer.job [-]: semantic enabled: clamping index_concurrency 6 -> 2 (memory bound: ...) +INFO indexer.job [-]: semantic enabled: clamping index_concurrency 6 -> 4 (memory bound: ...) ``` +At `index_concurrency <= 4` (the shipped default) the clamp is a no-op and this +line does not appear at all — confirmed on the live dev job at `index_concurrency: +4`, clamp 4: no clamp line, `4 worker(s)` on the disk line instead. + +**Re-derivation (#109).** A measured model — +`P_worst(N) = N * max((alpha+gamma)*B_breach, (alpha+gamma)*(d*C) + V_cap*C) + +extract_processes*R_proc + P_fixed`, coefficients measured across 4 corpora +(`alpha+gamma` up to 4.64 bytes-materialized per source byte, `d` up to 3748 +bytes/chunk, `V_cap` 40.1 KB/chunk resident), against a measured container +memory ceiling (`M`, read from cgroup + an allocate-until-failure bracket) — +showed N=4 clears 70% of `M` with a large margin at the standard 8000-chunk +global cap, and empirical confirmation on the real dev job agreed: two runs each +at N=2 and N=4 measured `peak rss: self=... children=...` (issue #109's new +instrumentation, `RUSAGE_SELF`/`RUSAGE_CHILDREN` at the end of `run()`) — N=2 +averaged ~90% of the 0.7*M budget, N=4 ~83%, both safely under, N=4 with *more* +margin. See `docs/perf/issue-109-measurements.md` for the full derivation, +every coefficient's provenance, and the two derived byte limits +(`MAX_EXTRACTED_BYTES` and the chunk cap `C`) this same model yields. + ### Embedding concurrency (#107) `workers x concurrency` is the number that matters, not `concurrency` alone. -Each of the (at most 2, semantic-clamped) workers dispatches up to -`semantic.embedding_concurrency` embedding batches at once +Each of the (at most 4, semantic-clamped -- #109 raised this from 2) workers +dispatches up to `semantic.embedding_concurrency` embedding batches at once (`app/embed.py:databricks_embedder`, order-preserving `ThreadPoolExecutor.map`): -2 x 4 = 8 in-flight gateway requests at the default, 2 x 8 = 16 at the -config.yaml-enforced ceiling of 8 (the `CODE_SEARCH_SEMANTIC_EMBEDDING_CONCURRENCY` -env var carries no ceiling, mirroring `semantic_embedding_batch_size`'s own -unbounded env surface -- config.yaml is the job's real surface regardless), both -under the SDK's 20-connection pool. +4 x 4 = 16 in-flight gateway requests at the default, **4 x 8 = 32 at the +config.yaml-enforced ceiling of 8 — this now EXCEEDS the SDK's 20-connection +pool** (the `CODE_SEARCH_SEMANTIC_EMBEDDING_CONCURRENCY` env var carries no +ceiling, mirroring `semantic_embedding_batch_size`'s own unbounded env surface +-- config.yaml is the job's real surface regardless). This combination was not +reachable before #109 (2 x 8 = 16 stayed under the pool); it is now, and is +flagged here rather than gated in code, matching this repo's "guardrail +constants with config-level fixes, not override flags" convention -- lower +`embedding_concurrency` if raising it alongside a near-ceiling +`index_concurrency`. `embedding_concurrency: 1` is the rollback switch — fully serial embedding, no thread pool spawned. See `docs/runbooks/semantic-enablement.md` §4 for the full in-flight/memory arithmetic and the 429 posture. @@ -408,7 +460,7 @@ or because the pool degraded. Three WARNING shapes to recognize: - **`... rebuilt the pool (generation N, rebuild M/3)`** — a worker died (`BrokenProcessPool`, e.g. a native crash in a grammar, an OOM kill). The branch(es) in flight at that moment failed (up to `index_concurrency` - branches, semantic-clamped to 2 — **not just one**: the pool is shared, so a + branches, semantic-clamped to 4 — **not just one**: the pool is shared, so a break can surface on every repo worker holding a future at that instant). Each failed branch re-indexes on its next run (it never got a stamp — the same self-healing property §1 describes). The pool is rebuilt and later @@ -450,7 +502,7 @@ retention (#105, §2.3) for the same in-flight branch. ### The connection pool follows the workers -Each worker holds exactly one connection, so the engine is built with +Each worker holds exactly one connection AT A TIME, so the engine is built with `pool_size == effective workers`, `max_overflow=0`, `pool_timeout=30`. There is deliberately **zero headroom**: a connection leak stalls loudly for 30 seconds and then raises, rather than growing the pool silently. If you see a @@ -458,6 +510,24 @@ and then raises, rather than growing the pool silently. If you see a connection, not an undersized pool — the pool is sized to the workers by construction. +**"One connection at a time" holds by sequencing, not by construction (#104, +verified by #109).** On the delta-gate-open semantic path a worker opens a +*second*, short-lived connection (`indexer/job.py`'s `with engine.connect() as +shas_conn:`, for the advisory `shas_fn` read) before its main `index_fn` +connection — but that second connection is closed before embedding/`index_fn` +starts, so at most one is ever open per worker at once. `pool_size == +effective_workers` (raised to 4 by #109 — see §3 above) stays correct only +because of that ordering; a future change that opened both connections +concurrently would silently under-provision the pool. Pinned by +`tests/unit/test_job.py::test_shas_fn_connection_closes_before_embedding_starts`. + +**Since #109 raised the semantic clamp from 2 to 4, the Lakebase connection +ceiling doubles on the semantic path too** — 4 concurrent connections at the +new default vs. 2 before. Not observed to be a bottleneck in either Arm B run +(no `QueuePool` timeouts, no degraded coverage), but worth knowing if a future +Lakebase compute-size change is on the table alongside a semantic-path +concurrency change. + The app/serving pool is separate and unaffected (5, paired with a matching `CapacityLimiter`). diff --git a/docs/runbooks/semantic-enablement.md b/docs/runbooks/semantic-enablement.md index e554afb..14e8032 100644 --- a/docs/runbooks/semantic-enablement.md +++ b/docs/runbooks/semantic-enablement.md @@ -81,7 +81,7 @@ semantic: ``` which makes the job a true semantic no-op (no embedder built, no chunking, the -2-worker clamp not applied) — no bundle/env change and no redeploy of the job's +4-worker clamp not applied) — no bundle/env change and no redeploy of the job's environment. Precedence for the job is `config.yaml > CODE_SEARCH_* env > default`, so `semantic.enabled: false` wins even if the env says enabled. @@ -141,45 +141,58 @@ buffered chunks, and that `semantic_max_chunks_per_repo` (`app/config.py`, defau silently truncating if a repo exceeds the ceiling. That default is deliberately conservative: the buffered vectors are Python float lists -costing ~32 B per element, so at `dim=1024` each chunk is ~32 KB and 8000 chunks is -~260 MB resident, held for the duration of the repo's write transaction. Raising it -scales memory linearly (50000 would be ~1.6 GB and would OOM a typical job container -*before* the loud ceiling check could fire, which defeats the purpose of the ceiling). -If a repo legitimately needs more, prefer the temp-table staging path (follow-up) over -raising this number. +costing ~32 B per element structural, but ~40.1 KB/chunk RESIDENT once measured (issue +#109 -- pymalloc overhead/fragmentation; use this figure for headroom arithmetic), so at +`dim=1024` 8000 chunks is ~313 MiB resident, held for the duration of the repo's write +transaction. Raising it scales memory roughly linearly (50000 would be ~1.9 GiB). Issue +#109 derived a full per-worker chunk-cap ceiling from the container-memory model: ~73,300 +chunks at the pinned N=2 semantic-worker count the derivation is evaluated at (methodology +only -- breaks a circularity in the model), ~36,700 at the shipped N=4 (see +`docs/perf/issue-109-measurements.md` §12) -- the current 8000 default uses well under a +quarter of either budget. A ceiling well past that derived bound (e.g. 50k) risks OOMing +the job container *before* the loud ceiling check could fire, which defeats the purpose of +the ceiling. If a repo legitimately needs more, prefer the temp-table staging path +(follow-up) over raising this number. **Per-repo override, without moving the global default:** `config.yaml`'s top-level `semantic_max_chunks_per_repo` map (`indexer/repo_config.py`) lets one outsized repo get its own cap — `indexer/resolve.py` carries the matched override onto that repo's `RepoEntry`, and `indexer/job.py` uses it in place of `cfg.semantic_max_chunks_per_repo` for that repo only (an active override is logged at INFO). It does not relax the -2-worker semantic clamp above, so a large override still multiplies whichever of the -(at most 2) concurrent workers happens to be indexing that repo — do the same ~32 -KB/chunk math against the override value, not just the global default, before setting -one. To move the **global** cap for the whole job instead of one repo, set +4-worker semantic clamp above (issue #109 raised this from 2), so a large override +still multiplies whichever of the (at most 4) concurrent workers happens to be +indexing that repo — do the same ~32 KB structural / ~40.1 KB resident per-chunk math +against the override value, not just the global default, before setting one. To move +the **global** cap for the whole job instead of one repo, set `semantic.max_chunks_per_repo` (a single int, section 6) — the map still wins for a repo it names. -**Parallelism:** with semantic on by default, the `effective_workers` clamp to 2 in -`indexer/job.py` now applies to every index run by default (each worker materialises a -whole repo's chunks) — see `docs/runbooks/indexing-parallelism.md`. +**Parallelism:** with semantic on by default, the `effective_workers` clamp to 4 +(issue #109; previously 2) in `indexer/job.py` now applies to every index run by +default (each worker materialises a whole repo's chunks) — see +`docs/runbooks/indexing-parallelism.md`. **Concurrent embedding requests (#107):** each worker's `embed()` call (`app/embed.py`) dispatches up to `semantic.embedding_concurrency` batches at once via a `ThreadPoolExecutor`, using `.map()` — never `as_completed()` — so vectors always come back in submission order regardless of which request finishes first. Total in-flight gateway requests for the job is `effective_workers x concurrency`: -2 x 4 = 8 at the default `embedding_concurrency: 4`, 2 x 8 = 16 at the config's -`le=8` ceiling, both under the SDK's 20-connection pool +4 x 4 = 16 at the default `embedding_concurrency: 4`, **4 x 8 = 32 at the config's +`le=8` ceiling — this now EXCEEDS the SDK's 20-connection pool** (issue #109 raised +`effective_workers` from 2 to 4; 2 x 8 = 16 stayed under the pool before) (`HTTPAdapter(pool_connections=20, pool_maxsize=20, pool_block=True)` — `pool_block=True` means exceeding the pool **silently serializes** requests rather -than raising, so staying under 20 is the load-bearing bound, not a nice-to-have). +than raising, so staying under 20 is the load-bearing bound, not a nice-to-have — a +run at the ceiling of both knobs still completes, just with less real concurrency +than configured). Lower `embedding_concurrency` if raising it alongside a +near-ceiling `index_concurrency`. The only new per-in-flight-batch memory cost is transient request/response buffers (~3.5 MB each: ~2.1 MB parsed vectors + ~1.3 MB raw JSON response + a -small request body) — ~28 MB at the default, ~56 MB at the ceiling, negligible -beside the ~0.5–0.8 GB/worker baseline above. `embedding_concurrency: 1` restores -today's fully serial embedding and spawns no thread pool at all (the rollback -switch). +small request body) — ~56 MB at the default (4 workers x 4), ~112 MB at the +ceiling, negligible beside the ~313 MiB/worker vector baseline above (8000 +chunks at the resident 40.1 KB/chunk figure). `embedding_concurrency: 1` +restores today's fully serial embedding and spawns no thread pool at all (the +rollback switch). 429s from the AI Gateway are absorbed entirely by the `databricks-sdk`'s own `Retry-After`-honouring backoff (`_RetryAfterCustomizer`, defaulting to 1s when diff --git a/indexer/ingest.py b/indexer/ingest.py index e630cf3..02737aa 100644 --- a/indexer/ingest.py +++ b/indexer/ingest.py @@ -126,6 +126,18 @@ # bounds members of ANY type. Regular-file accounting alone would let an # archive of a million directory or link headers decompress unbounded -- # they carry no data, so `streamed` never moves. +# +# issue #109 re-derived this as a MEMORY limit (previously undirected): at the +# pinned N=2 operating point, against a measured container ceiling and the +# semantic path's measured bytes-materialized-per-source-byte coefficients, the +# breach-regime peak stays under 0.7x that ceiling only for a per-branch source +# total <= ~880 MiB (vs this 2 GB). NOT changed here: no branch in the +# measurement corpus approached either value (largest observed: ~31.5 MiB, ~3.6% +# of the derived bound), so there is no empirical signal to justify lowering a +# constant whose breach fails the branch and closes the whole run's +# reconciliation checkpoint (job.py's `_decide_reconciliation` requires +# `failures == 0`). See docs/perf/issue-109-measurements.md for the full +# derivation. MAX_EXTRACTED_BYTES = 2_000_000_000 diff --git a/indexer/job.py b/indexer/job.py index 16e8d23..f7631ce 100644 --- a/indexer/job.py +++ b/indexer/job.py @@ -121,6 +121,7 @@ import argparse import base64 import logging +import resource import shutil import sys import tempfile @@ -410,10 +411,16 @@ def run( owns_engine = engine is None if engine is None: # The pool is DERIVED from the worker count, not a constant: each worker - # holds exactly one connection (one engine.connect() per repo, and the - # embed/chunk precompute happens before it), so pool_size == workers is - # exactly enough and max_overflow=0 turns a connection leak into a loud - # stall instead of silent pool growth. pool_timeout is SQLAlchemy's own + # holds at most ONE connection AT A TIME -- by sequencing, not by + # construction (issue #109 E7). On the delta-gate-open semantic path a + # worker opens a second, short-lived connection for the advisory + # shas_fn read (below), but closes it before the embed/chunk precompute + # and its own index_fn connection open -- see + # tests/unit/test_job.py::test_shas_fn_connection_closes_before_embedding_starts. + # So pool_size == workers is exactly enough and max_overflow=0 turns a + # connection leak into a loud stall instead of silent pool growth (a + # future change that opened both connections concurrently would need a + # bigger pool). pool_timeout is SQLAlchemy's own # default, spelled out HERE because max_overflow=0 is what makes it # observable -- a reader seeing the overflow ban must not have to go look # up how long the resulting stall lasts. Passing pool_size explicitly is @@ -616,6 +623,16 @@ def run( len(entries), time.monotonic() - run_started, ) + # issue #109 AC3: peak RSS for this run, self (the main process, every repo + # worker thread) and children (the #108 extraction pool's worker processes, + # otherwise invisible from here) -- ru_maxrss is a high-water mark, in KB on + # Linux, so no delta/baseline subtraction is needed the way a local + # measurement script needs one against its own import-time baseline. + logger.info( + "peak rss: self=%d KB children=%d KB", + resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss, + ) # Conflicts do NOT fail the run because they SELF-HEAL -- the stamp that # displaced them makes the next run re-index that branch unconditionally. # Note this trades a paging signal for one run of staleness on that branch; diff --git a/indexer/repo_config.py b/indexer/repo_config.py index 38cacf5..fb42c7c 100644 --- a/indexer/repo_config.py +++ b/indexer/repo_config.py @@ -270,16 +270,25 @@ class RepoConfig(BaseModel): from ``index_concurrency`` entirely. Raise ``index_concurrency`` only knowing you are buying disk-bound repo fan-out, not extraction throughput. - When semantic indexing is on, the effective worker count is clamped to 2 by - :func:`effective_workers`. That clamp is a **memory** bound, not a CPU one: - embedding materialises a whole repo's chunks in memory (~0.5-0.8 GB per - worker; 260 MB of vectors alone at the 8000-chunk ceiling). + When semantic indexing is on, the effective worker count is clamped to 4 by + :func:`effective_workers` (issue #109 re-derived this from 2: a measured + ``P_worst`` model -- ``(alpha+gamma)`` bytes-materialized-per-source-byte + coefficients, measured resident vector cost, #108's per-process RSS, and a + measured container memory ceiling -- showed N=4 clears 0.7x the container + budget with margin, and two live-job runs at N=4 confirmed it empirically: + peak self+children RSS landed at ~83% of budget, actually MORE comfortable + than N=2's own ~90%. See docs/perf/issue-109-measurements.md). That clamp is + still a **memory** bound, not a CPU one: embedding materialises a whole + repo's chunks in memory. Per-chunk vector cost is ~32 KB structural (dim=1024 + Python float-list storage) but ~40.1 KB RESIDENT (measured; pymalloc + overhead/fragmentation) -- use the resident figure for headroom arithmetic -- + so 313 MiB of vectors alone at the 8000-chunk ceiling. ``semantic_max_chunks_per_repo`` (the per-repo MAP) overrides that global 8000-chunk ceiling for individual repos named here, without moving the global - default. It does NOT relax the 2-worker semantic clamp above -- a large + default. It does NOT relax the 4-worker semantic clamp above -- a large override still multiplies the per-worker memory cost of whichever of the (at - most 2) concurrent semantic workers happens to be indexing that repo. + most 4) concurrent semantic workers happens to be indexing that repo. The similarly-named ``semantic.max_chunks_per_repo`` (inside the ``semantic:`` block, a single INT) moves that GLOBAL ceiling itself for the whole job. The @@ -354,11 +363,17 @@ def _normalize_semantic_overrides(self) -> RepoConfig: def effective_workers(config: RepoConfig, *, semantic_enabled: bool) -> int: """Worker-pool size for a run, applying the semantic memory clamp. + The clamp is 4 (issue #109; previously 2 -- see :class:`RepoConfig`'s + docstring for the re-derivation and empirical Arm A/B confirmation). It is a + ceiling, never a floor: an ``index_concurrency`` below the clamp passes + through unchanged, so N=3 (or any other legal value) is a real, reachable + ``effective_workers`` outcome, not just N in {1, 2, 4}. + Takes a plain ``bool`` rather than ``Settings`` so this module keeps its import-light property (see the module docstring). """ if semantic_enabled: - return min(config.index_concurrency, 2) + return min(config.index_concurrency, 4) return config.index_concurrency diff --git a/scripts/measure_ingest_threads.py b/scripts/measure_ingest_threads.py new file mode 100644 index 0000000..8a0e23b --- /dev/null +++ b/scripts/measure_ingest_threads.py @@ -0,0 +1,473 @@ +"""Measure `iter_tar_source_files` thread-scaling and its GIL-bound component +mix (#109 §3.4, E4 -- "M2: does the ingest pass scale across threads?"). + +Since #106 the per-branch file source is a single serial pass over one open +``TarFile`` (:func:`indexer.ingest.iter_tar_source_files`): a ``tf.next()`` +member walk, ``fh.read()``, the NUL-sniff binary check +(:func:`indexer.parse._looks_binary`), a UTF-8 decode, and a NUL-strip. Raising +``index_concurrency`` multiplies *concurrent* ingest passes across repo-worker +threads, and whether that helps is gated on which of those steps hold the GIL. +A planning-time synthetic probe (40 MB payload, isolated ``bytes.decode`` vs +``zlib.decompress``) found decode GIL-bound and zlib GIL-releasing; this script +re-measures the FULL real pass, decomposed, over real tarballs, rather than +re-quoting that probe. + +Two things are measured, and the "decompose by component" and "end-to-end +speedup" numbers deliberately come from two different code paths per the +plan's own instruction: + +* **Component breakdown** -- a re-walk of the tarball using the SAME private + primitives ``iter_tar_source_files`` calls (``indexer.ingest + ._normalise_member_name`` / ``._assert_link_target_is_contained``, + ``indexer.parse._looks_binary``), imported rather than re-forked, with a + ``time.perf_counter()`` pair around each step. This is a timing-only + reimplementation of the loop -- never the source of the speedup numbers. +* **End-to-end speedup** -- always calls the REAL + ``indexer.ingest.iter_tar_source_files`` and nothing else, N threads each + streaming its OWN distinct real tarball (page cache and content mix must not + be shared across threads), compared against a single-thread sequential pass + over the SAME N tarballs. + +Two conditions: CPU otherwise idle, and with :class:`indexer.extract_pool. +ExtractionPool` constructed and continuously driving real ``.stream()`` +extraction on a background thread, to see whether process-pool contention for +cores changes the picture. + +**The tarballs must be real** -- fetched over HTTP from public GitHub repos via +:func:`indexer.fetch.download_tarball` (unauthenticated ``httpx.Client`` works +fine for public repos) and cached locally so a re-run doesn't re-download. + +Usage: ``uv run python scripts/measure_ingest_threads.py [--repeat 3] +[--pool-processes 4] [--cache-dir /tmp/measure_ingest_threads_cache] +[--repos org/repo@ref,org/repo@ref,...]`` (need >= 4 distinct repos). +""" + +from __future__ import annotations + +import argparse +import gzip +import tarfile +import threading +import time +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import TypeVar + +import httpx + +from indexer.extract_pool import ExtractionPool +from indexer.fetch import download_tarball +from indexer.ingest import ( + MAX_EXTRACTED_BYTES, + _assert_link_target_is_contained, + _normalise_member_name, + iter_tar_source_files, +) +from indexer.languages import MAX_FILE_BYTES +from indexer.parse import _looks_binary + +T = TypeVar("T") + +# Four distinct public repos of comparable decoded size (~20-37 MB each, +# verified by hand before picking these four) -- comparable size matters for +# the N=4 comparison specifically: four wildly mismatched tarballs would let +# the largest one dominate both the sequential sum and the concurrent wall +# clock, measuring "how fast is the biggest repo alone" rather than genuine +# 4-way thread scaling. +DEFAULT_REPOS = [ + ("tiangolo", "fastapi", "master"), + ("sqlalchemy", "sqlalchemy", "main"), + ("sympy", "sympy", "master"), + ("django", "django", "main"), +] + + +def _parse_repos(spec: str) -> list[tuple[str, str, str]]: + out = [] + for item in spec.split(","): + org_repo, _, ref = item.partition("@") + org, _, repo = org_repo.partition("/") + out.append((org, repo, ref or "HEAD")) + return out + + +def _cache_path(cache_dir: Path, org: str, repo: str) -> Path: + return cache_dir / f"{org}__{repo}.tar.gz" + + +def fetch_tarballs(repos: Sequence[tuple[str, str, str]], cache_dir: Path) -> list[Path]: + """Download each repo's tarball once (cached under ``cache_dir`` across runs).""" + cache_dir.mkdir(parents=True, exist_ok=True) + paths = [] + with httpx.Client(timeout=120.0) as client: + for org, repo, ref in repos: + dest = _cache_path(cache_dir, org, repo) + if not dest.exists(): + print(f"fetching {org}/{repo}@{ref} ...") + tmp_dir = cache_dir / f"_dl_{org}_{repo}" + downloaded = download_tarball(client, org, repo, ref, tmp_dir) + downloaded.replace(dest) + try: + tmp_dir.rmdir() + except OSError: + pass + paths.append(dest) + return paths + + +@dataclass +class ComponentTimes: + """Cumulative wall-clock seconds per component of one (or more, summed) + instrumented walks -- see :func:`decompose_walk`.""" + + tf_next: float = 0.0 + fh_read: float = 0.0 + looks_binary: float = 0.0 + decode: float = 0.0 + nul_strip: float = 0.0 + n_files: int = 0 + n_members: int = 0 + + def add(self, other: "ComponentTimes") -> None: + self.tf_next += other.tf_next + self.fh_read += other.fh_read + self.looks_binary += other.looks_binary + self.decode += other.decode + self.nul_strip += other.nul_strip + self.n_files += other.n_files + self.n_members += other.n_members + + @property + def total(self) -> float: + return self.tf_next + self.fh_read + self.looks_binary + self.decode + self.nul_strip + + +def decompose_walk(tar_path: Path) -> ComponentTimes: + """Re-walk ``tar_path`` with the exact filter chain + ``indexer.ingest.iter_tar_source_files`` uses, timing each component with + its own ``perf_counter()`` pair. + + A TIMING-ONLY reimplementation of that function's loop body: it imports the + same private primitives (``_normalise_member_name``, + ``_assert_link_target_is_contained``, ``_looks_binary``) rather than + re-forking their logic, so the filter population -- which members get as + far as ``fh.read()`` / decode -- matches production exactly. It does not + build ``ParsedFile`` objects or return content; the real end-to-end number + always comes from calling ``iter_tar_source_files`` itself (see + :func:`run_sequential_real` / :func:`run_concurrent_real`), never from this + function. + """ + ct = ComponentTimes() + tf = tarfile.open(tar_path, mode="r:gz") + try: + top_dir: str | None = None + streamed = 0 + seen: set[str] = set() + while True: + t0 = time.perf_counter() + member = tf.next() + ct.tf_next += time.perf_counter() - t0 + if member is None: + break + ct.n_members += 1 + tf.members.clear() # type: ignore[attr-defined] + + name = _normalise_member_name(member.name) + component = name.split("/", 1)[0] + if top_dir is None: + top_dir = component + elif component != top_dir: + raise ValueError( + "expected exactly one top-level dir in tarball, found " + f"{sorted({top_dir, component})}" + ) + + if member.islnk() or member.issym(): + _assert_link_target_is_contained(name, member.linkname) + + if member.offset > MAX_EXTRACTED_BYTES: + raise ValueError( + f"tarball stream reaches {member.offset} decompressed bytes, " + f"exceeding {MAX_EXTRACTED_BYTES}" + ) + + if not member.isreg(): + continue + + streamed += member.size + if streamed > MAX_EXTRACTED_BYTES: + raise ValueError( + f"tarball streams to {streamed} bytes of content, " + f"exceeding {MAX_EXTRACTED_BYTES}" + ) + + rel_path = name[len(top_dir) + 1 :] if name != top_dir else "" + if not rel_path: + continue + if ".git" in rel_path.split("/"): + continue + if member.size > MAX_FILE_BYTES: + continue + if rel_path in seen: + continue + seen.add(rel_path) + + fh = tf.extractfile(member) + if fh is None: + continue + t0 = time.perf_counter() + with fh: + raw = fh.read() + ct.fh_read += time.perf_counter() - t0 + + t0 = time.perf_counter() + is_binary = _looks_binary(raw) + ct.looks_binary += time.perf_counter() - t0 + if is_binary: + continue + + t0 = time.perf_counter() + try: + content = raw.decode("utf-8") + except UnicodeDecodeError: + ct.decode += time.perf_counter() - t0 + continue + ct.decode += time.perf_counter() - t0 + + t0 = time.perf_counter() + content.replace("\x00", "") + ct.nul_strip += time.perf_counter() - t0 + ct.n_files += 1 + finally: + tf.close() + return ct + + +def measure_zlib_isolated(tar_path: Path) -> float: + """Whole-archive gzip inflate, isolated from tarfile header parsing -- + the cleanest available measurement of the GIL-releasing component named in + the plan's §0.4 prior (``zlib.decompress``, 2.54x at 4 threads on a 40 MB + synthetic payload).""" + raw_gz = tar_path.read_bytes() + start = time.perf_counter() + gzip.decompress(raw_gz) + return time.perf_counter() - start + + +def best_of(fn: Callable[[], T], repeat: int, key: Callable[[T], float]) -> T: + """>=3 repeats, discard one warm-up, report best-of -- the methodology + ``docs/perf/issue-108-measurements.md`` used for #108.""" + results = [fn() for _ in range(max(repeat, 1))] + kept = results[1:] if len(results) > 1 else results + return min(kept, key=key) + + +def run_sequential_real(paths: Sequence[Path]) -> float: + """Single thread, ``iter_tar_source_files`` over every path in ``paths``, + one after another -- the baseline half of the fair sequential-vs-concurrent + comparison (same tarball set both sides).""" + start = time.perf_counter() + for p in paths: + list(iter_tar_source_files(p)) + return time.perf_counter() - start + + +def run_concurrent_real(paths: Sequence[Path]) -> float: + """``len(paths)`` threads, each streaming its OWN tarball via the REAL + ``iter_tar_source_files`` concurrently.""" + start = time.perf_counter() + with ThreadPoolExecutor(max_workers=len(paths)) as ex: + futures = [ex.submit(lambda p=p: list(iter_tar_source_files(p))) for p in paths] + for f in futures: + f.result() + return time.perf_counter() - start + + +def run_concurrent_decomp(paths: Sequence[Path]) -> tuple[float, ComponentTimes]: + """``len(paths)`` threads, each running the instrumented + :func:`decompose_walk` on its own tarball concurrently. Returns the + wall-clock for the whole concurrent run plus the SUM of every thread's + per-component cumulative time, so "does component X's aggregate cost grow + linearly with N" is directly readable off two consecutive rows.""" + start = time.perf_counter() + with ThreadPoolExecutor(max_workers=len(paths)) as ex: + results = [f.result() for f in [ex.submit(decompose_walk, p) for p in paths]] + elapsed = time.perf_counter() - start + total = ComponentTimes() + for r in results: + total.add(r) + return elapsed, total + + +class PoolContention: + """Drives ``ExtractionPool.stream()`` continuously on a background thread + over a fixed file list, to create real multi-process CPU contention for + the "pool live" condition. Started/stopped once per condition, spanning + every N in that condition's matrix.""" + + def __init__(self, files: list, n_processes: int) -> None: + self._pool = ExtractionPool(n_processes=n_processes) + self._files = files + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def _run(self) -> None: + while not self._stop.is_set(): + list(self._pool.stream(self._files)) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=30) + self._pool.shutdown() + + +def _bucket(n: int, speedup: float) -> str: + """Fixed-in-advance thresholds from the plan's §3.4.""" + if n == 2: + if speedup >= 1.6: + return "PARALLELIZES" + if speedup <= 1.1: + return "NO-SCALE" + return "ambiguous" + if n == 4: + if speedup >= 2.0: + return "PARALLELIZES" + if speedup <= 1.2: + return "NO-SCALE" + return "PARTIAL" + return "n/a" + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cache-dir", type=Path, default=Path("/tmp/measure_ingest_threads_cache")) + parser.add_argument("--repeat", type=int, default=3, help=">=3 per the plan's protocol") + parser.add_argument( + "--pool-processes", + type=int, + default=4, + help="ExtractionPool size for the 'pool live' condition", + ) + parser.add_argument( + "--repos", + default=",".join(f"{o}/{r}@{ref}" for o, r, ref in DEFAULT_REPOS), + help="comma-separated org/repo@ref list, need >= 4 distinct repos", + ) + args = parser.parse_args(argv) + + repos = _parse_repos(args.repos) + if len(repos) < 4: + raise SystemExit( + "need >= 4 distinct repos for the N in {1..4} thread-scaling matrix, " + f"got {len(repos)}" + ) + + paths = fetch_tarballs(repos, args.cache_dir) + + print("=" * 100) + print(f"tarballs (cached under {args.cache_dir}):") + for (org, repo, ref), p in zip(repos, paths, strict=True): + files = list(iter_tar_source_files(p)) + total_bytes = sum(len(pf.content) for pf in files) + print( + f" {org}/{repo}@{ref}: {p.name}, {p.stat().st_size / 1e6:.2f} MB compressed, " + f"{len(files)} indexable files, {total_bytes / 1e6:.2f} MB decoded text" + ) + print("=" * 100) + + # ---- Part 1: per-tarball component breakdown, single-threaded, idle CPU ---- + print("\n### Part 1 -- per-tarball component breakdown (single-threaded, CPU idle) ###") + print("component times are cumulative seconds inside the instrumented re-walk (see docstring);") + print( + "zlib_iso_s is a SEPARATE standalone whole-archive gzip.decompress(), not part of the " + "walk sum.\n" + ) + print( + f"{'repo':>22s} {'zlib_iso_s':>10s} {'tf_next_s':>10s} {'fh_read_s':>10s} " + f"{'binary_s':>9s} {'decode_s':>9s} {'nulstrip_s':>10s} {'n_files':>7s} " + f"{'n_members':>9s}" + ) + for (org, repo, ref), p in zip(repos, paths, strict=True): + zlib_s = best_of(lambda p=p: measure_zlib_isolated(p), args.repeat, key=lambda x: x) + ct = best_of(lambda p=p: decompose_walk(p), args.repeat, key=lambda c: c.total) + label = f"{org}/{repo}" + print( + f"{label:>22s} {zlib_s:>10.4f} {ct.tf_next:>10.4f} {ct.fh_read:>10.4f} " + f"{ct.looks_binary:>9.4f} {ct.decode:>9.4f} {ct.nul_strip:>10.4f} " + f"{ct.n_files:>7d} {ct.n_members:>9d}" + ) + + # ---- Part 2: N-thread scaling matrix, two conditions ---- + for condition in ("idle", "pool_live"): + print(f"\n### Part 2 -- N-thread ingest scaling, condition={condition} ###") + contention: PoolContention | None = None + if condition == "pool_live": + load_files = list(iter_tar_source_files(paths[0])) + print( + f"starting background ExtractionPool(n_processes={args.pool_processes}) " + f"driving .stream() over {len(load_files)} files from " + f"{repos[0][0]}/{repos[0][1]} ..." + ) + contention = PoolContention(load_files, n_processes=args.pool_processes) + contention.start() + time.sleep(0.5) # let the process pool spin up before timing starts + try: + print( + "\nend-to-end speedup (REAL iter_tar_source_files; same N-tarball set both sides):" + ) + print( + f"{'N':>3s} {'seq_s (1thr, Nx)':>17s} {'conc_s (Nthr)':>14s} " + f"{'speedup':>9s} {'bucket':>13s}" + ) + for n in (1, 2, 3, 4): + subset = paths[:n] + seq = best_of( + lambda subset=subset: run_sequential_real(subset), args.repeat, key=lambda x: x + ) + if n > 1: + conc = best_of( + lambda subset=subset: run_concurrent_real(subset), + args.repeat, + key=lambda x: x, + ) + else: + conc = seq + speedup = seq / conc if conc else float("inf") + bucket = _bucket(n, speedup) + print(f"{n:>3d} {seq:>17.4f} {conc:>14.4f} {speedup:>8.2f}x {bucket:>13s}") + + print( + "\ncomponent breakdown at each N (instrumented re-walk, N threads concurrent, " + "SUM across threads):" + ) + print( + f"{'N':>3s} {'wall_s':>8s} {'sum_tf_next_s':>13s} {'sum_fh_read_s':>13s} " + f"{'sum_binary_s':>12s} {'sum_decode_s':>12s} {'sum_nulstrip_s':>14s}" + ) + for n in (1, 2, 3, 4): + subset = paths[:n] + elapsed, ct = best_of( + lambda subset=subset: run_concurrent_decomp(subset), + args.repeat, + key=lambda t: t[0], + ) + print( + f"{n:>3d} {elapsed:>8.4f} {ct.tf_next:>13.4f} {ct.fh_read:>13.4f} " + f"{ct.looks_binary:>12.4f} {ct.decode:>12.4f} {ct.nul_strip:>14.4f}" + ) + finally: + if contention is not None: + contention.stop() + print("stopped background ExtractionPool contention.") + + print("\ndone.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/measure_semantic_memory.py b/scripts/measure_semantic_memory.py new file mode 100644 index 0000000..a57d397 --- /dev/null +++ b/scripts/measure_semantic_memory.py @@ -0,0 +1,574 @@ +"""Measure the semantic path's peak memory, decomposed into named terms (#109, AC3). + +Offline companion to ``docs/perf/issue-109-measurements.md`` §3.3(b)'s ``P_worst`` +model. Drives the REAL production call sequence from +``indexer.job._index_one_branch``'s semantic path -- never a re-implementation: + + files = list(iter_tar_source_files(tar_path)) # alpha + files_to_embed = [pf for pf in files # the delta gate + if (pf.path, content_sha(pf.content)) not in carried] + chunk_writer = indexer.job._precompute_chunk_writer( # gamma + vectors + files_to_embed, embed_fn, max_chunks_per_repo) + +``_precompute_chunk_writer`` is private (leading underscore) but imported directly +by name -- this is a measurement script living in the same repo, not an external +consumer, and the alternative (re-typing its chunking/embedding logic here) is +exactly the drift this script exists to avoid. + +**The stub embedder returns DISTINCT floats per vector component** +(``float(i * dim + j)``), never a shared/cached float such as ``[0.0] * dim``. +Planning found that the naive stub understates resident memory by ~4x, because +``[0.0] * dim`` stores ``dim`` references to ONE cached float object rather than +``dim`` independent ``PyFloat`` allocations (see the plan's §2.2). This script's +correctness as a memory probe depends entirely on avoiding that trap. + +**Methodology -- three named RSS terms per corpus/gate combination:** + +Peak RSS (``resource.getrusage(RUSAGE_SELF).ru_maxrss``) is a monotonic +non-decreasing high-water mark *within one process*, so every stage below runs in +its OWN freshly spawned subprocess (this same script, re-invoked with a hidden +``--worker`` flag) -- otherwise stage N's baseline would already carry stage +N-1's peak forward and every delta after the first would be contaminated. Inside +one subprocess, three readings bound three terms: + + 1. baseline (interpreter only) + 2. after ``files = list(iter_tar_source_files(tar))`` -> the FILES term (alpha) + 3. after ``per_file = {p: list(iter_chunks(p)) for p in files_to_embed}`` + -> the CHUNKS term (gamma) + 4. after ``_precompute_chunk_writer(files_to_embed, stub_embed_fn, huge_cap)`` + -> the VECTORS term + +Between readings 3 and 4 the manually-built ``per_file`` dict from step 3 is +deleted and ``gc.collect()``ed *before* calling the real ``_precompute_chunk_writer``, +which re-chunks internally (it has no way to accept precomputed chunks -- that +would be re-implementing its contract, not measuring it). The intent is that the +freed step-3 allocations are reused by the allocator for step 4's structurally +identical rebuild, so the reading-3-to-4 delta is dominated by the NEW allocation +(the embedding vectors) rather than by double-counting the chunk objects. This is +a reasonable approximation on CPython/glibc for same-shaped allocations, not a +guarantee -- reported numbers may run slightly high for exactly this reason, and +that is called out again in the printed report. + +The chunk cap passed to ``_precompute_chunk_writer`` here is deliberately huge +(never the production ``semantic_max_chunks_per_repo``): this script measures the +UNCAPPED terms (alpha, gamma) and the per-chunk vector cost directly, not +cap-breach behavior -- that is a different, later step of issue #109's plan +(§3.3(c)-(f)), not this script's job. + +**V_cap (resident bytes per chunk)** is measured completely separately from the +corpus runs, matching the plan's §2.2 methodology exactly: build exactly 8000 +distinct-float 1024-dim vectors directly (not through ``_precompute_chunk_writer``) +and take one baseline/after delta. + +**N-concurrency (N in {1,2,3,4})** spins up N ``threading.Thread``s, each running +the SAME files -> chunk_writer pipeline against a real tarball, with ONE shared +``indexer.extract_pool.ExtractionPool`` built up front and each thread draining +its own ``pool.stream(files)`` call to completion -- so the #108 process pool's +own resident overhead (R_proc) is live and contributing to the measured peak, +exactly like N concurrent repo-worker threads in production. Each thread builds +its own ``files`` list from its own call to ``iter_tar_source_files`` (never a +shared generator -- the module's own docstring warns that a second thread +advancing the same tar-backed generator corrupts output silently). Both +``RUSAGE_SELF`` (this process, all threads) and ``RUSAGE_CHILDREN`` (the pool's +worker processes) are reported. + +**Corpora**: this repo (``IceRhymers/databricks-code-search``) at HEAD, plus three +other modest real public repos spanning a size range (Flask, Requests, Django), +fetched via the real ``indexer.fetch.download_tarball`` and cached under +``--cache-dir`` so repeat runs do not re-download. Total download is a few tens of +MB, well under the ~200 MB budget. + +Usage: ``uv run python scripts/measure_semantic_memory.py`` +(no arguments needed for the default corpus set; see ``--help`` for overrides). +""" + +from __future__ import annotations + +import argparse +import gc +import itertools +import json +import logging +import random +import resource +import subprocess +import sys +import threading +import time +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +import httpx + +from indexer.extract_pool import ExtractionPool, _available_cpus +from indexer.fetch import download_tarball +from indexer.hashing import content_sha +from indexer.ingest import iter_tar_source_files +from indexer.languages import ParsedFile +from indexer.parse import iter_chunks + +logging.disable(logging.CRITICAL) # keep worker-subprocess stdout free of log noise + +# A generous cap -- never the production semantic_max_chunks_per_repo -- so +# _precompute_chunk_writer never raises the cap-breach ValueError while this +# script measures the uncapped terms it deliberately does not exercise here. +_UNCAPPED_MAX_CHUNKS_PER_REPO = 50_000_000 + +# Default corpus: this repo plus three modest, well-known public repos spanning +# a size range. `ref` is always "HEAD" -- download_tarball's `ref` argument is +# passed straight into GitHub's tarball URL and accepts a branch name. +_DEFAULT_CORPORA = [ + ("databricks-code-search", "IceRhymers", "databricks-code-search"), + ("flask", "pallets", "flask"), + ("requests", "psf", "requests"), + ("django", "django", "django"), +] + +_GATE_STATES: list[tuple[str, float]] = [ + ("first-index", 0.0), # gate closed: files_to_embed IS files, no narrowing + ("recurring-1pct", 0.01), # gate open: carried covers ~99% of files + ("recurring-10pct", 0.10), # gate open: carried covers ~90% of files +] + + +def _rss_kb() -> int: + """Peak RSS of THIS process so far, in KB (Linux ``ru_maxrss`` semantics).""" + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + + +def _children_rss_kb() -> int: + """Peak RSS across reaped child processes, in KB.""" + return resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss + + +def _stub_embed_fn(dim: int = 1024): + """A stub ``EmbedFn`` returning DISTINCT floats per vector component. + + NEVER ``[0.0] * dim`` -- that stores ``dim`` references to one cached float + and understates resident memory by ~4x (planning's §2.2 trap). A running + counter guarantees every float, across every call, is a fresh ``PyFloat``. + """ + counter = itertools.count() + + def embed_fn(texts: list[str]) -> list[list[float]]: + vectors = [] + for _ in texts: + base = next(counter) * dim + vectors.append([float(base + j) for j in range(dim)]) + return vectors + + return embed_fn + + +def _source_bytes(files: Sequence[ParsedFile]) -> int: + """Total UTF-8-encoded byte length of ``files``' content -- the denominator + for both alpha and gamma.""" + return sum(len(pf.content.encode("utf-8")) for pf in files) + + +def _narrow_files_to_embed( + files: list[ParsedFile], gate: str, delta_fraction: float, *, seed: int = 0 +) -> list[ParsedFile]: + """Replicate ``_index_one_branch``'s exact narrowing logic for a simulated + gate state. + + ``gate == "first-index"``: the gate is CLOSED (no stamp at the current + ``INDEX_SEMANTICS_VERSION``), so production sets ``files_to_embed = files`` + verbatim -- no ``carried`` set is even read. This is the shape that OOMs + (§0.2 of the plan): no narrowing benefit at all. + + Otherwise: the gate is OPEN, and ``carried`` is simulated as covering + ``1 - delta_fraction`` of ``files``' ``(path, content_sha)`` pairs (a + deterministic random sample), so ``files_to_embed`` narrows to approximately + ``delta_fraction`` of ``files`` -- using the SAME list-comprehension shape + ``indexer/job.py`` uses, not a re-derived equivalent. + """ + if gate == "first-index": + return files + shas = [(pf.path, content_sha(pf.content)) for pf in files] + n_carry = round(len(shas) * (1 - delta_fraction)) + rng = random.Random(seed) + carried = set(rng.sample(shas, n_carry)) if shas else set() + return [pf for pf in files if (pf.path, content_sha(pf.content)) not in carried] + + +# -------------------------------------------------------------------------- +# Worker bodies -- each runs in ITS OWN freshly spawned subprocess (see the +# module docstring for why ru_maxrss's high-water-mark semantics demand this). +# -------------------------------------------------------------------------- + + +def _worker_stage(cfg: dict[str, Any]) -> dict[str, Any]: + from indexer.job import _precompute_chunk_writer # local: keep worker startup lean + + tarball = Path(cfg["tarball"]) + gate = cfg["gate"] + delta_fraction = cfg["delta_fraction"] + + baseline_kb = _rss_kb() + + files = list(iter_tar_source_files(tarball)) + after_files_kb = _rss_kb() + + files_to_embed = _narrow_files_to_embed(files, gate, delta_fraction) + + per_file_chunks = {pf.path: list(iter_chunks(pf)) for pf in files_to_embed} + after_chunks_kb = _rss_kb() + chunk_count = sum(len(chunks) for chunks in per_file_chunks.values()) + chunk_content_bytes = sum( + len(c.content.encode("utf-8")) for chunks in per_file_chunks.values() for c in chunks + ) + del per_file_chunks + gc.collect() + + embed_fn = _stub_embed_fn() + chunk_writer = _precompute_chunk_writer(files_to_embed, embed_fn, _UNCAPPED_MAX_CHUNKS_PER_REPO) + after_vectors_kb = _rss_kb() + del chunk_writer + + return { + "baseline_kb": baseline_kb, + "after_files_kb": after_files_kb, + "after_chunks_kb": after_chunks_kb, + "after_vectors_kb": after_vectors_kb, + "files_count": len(files), + "files_bytes": _source_bytes(files), + "files_to_embed_count": len(files_to_embed), + "files_to_embed_bytes": _source_bytes(files_to_embed), + "chunk_count": chunk_count, + "chunk_content_bytes": chunk_content_bytes, + } + + +def _worker_vcap(cfg: dict[str, Any]) -> dict[str, Any]: + n = cfg["n"] + dim = cfg["dim"] + baseline_kb = _rss_kb() + vectors = [[float(i * dim + j) for j in range(dim)] for i in range(n)] + after_kb = _rss_kb() + assert len(vectors) == n and len(vectors[0]) == dim + return {"baseline_kb": baseline_kb, "after_kb": after_kb, "n": n, "dim": dim} + + +def _worker_nconc(cfg: dict[str, Any]) -> dict[str, Any]: + from indexer.job import _precompute_chunk_writer # local: keep worker startup lean + + tarball = Path(cfg["tarball"]) + n_threads = cfg["n_threads"] + gate = cfg["gate"] + delta_fraction = cfg["delta_fraction"] + + baseline_self_kb = _rss_kb() + baseline_children_kb = _children_rss_kb() + + n_processes = min(_available_cpus(), 8) + pool = ExtractionPool(n_processes=n_processes) + + errors: list[str] = [] + + def run_one() -> None: + try: + files = list(iter_tar_source_files(tarball)) + files_to_embed = _narrow_files_to_embed(files, gate, delta_fraction) + embed_fn = _stub_embed_fn() + # Held alive (not discarded) across pool.stream() below: production + # (indexer/job.py) keeps chunk_writer's vectors resident for the whole + # index_repo write window, which is exactly the concurrent-residency + # property this stage measures -- freeing it early would let each + # thread's vectors be collected before, or concurrently with, sibling + # threads' peaks, understating true N-way concurrent RSS. + chunk_writer = _precompute_chunk_writer( + files_to_embed, embed_fn, _UNCAPPED_MAX_CHUNKS_PER_REPO + ) + # Drain the pool's stream fully so its worker processes actually do + # (and stay resident for) the same work a real branch would ask of them. + list(pool.stream(files)) + del chunk_writer + except Exception as exc: # noqa: BLE001 -- reported, not swallowed + errors.append(repr(exc)) + + threads = [threading.Thread(target=run_one) for _ in range(n_threads)] + start = time.perf_counter() + for t in threads: + t.start() + for t in threads: + t.join() + elapsed = time.perf_counter() - start + pool.shutdown() + + after_self_kb = _rss_kb() + after_children_kb = _children_rss_kb() + + return { + "n_threads": n_threads, + "n_processes": n_processes, + "baseline_self_kb": baseline_self_kb, + "after_self_kb": after_self_kb, + "baseline_children_kb": baseline_children_kb, + "after_children_kb": after_children_kb, + "elapsed_s": elapsed, + "errors": errors, + } + + +_WORKERS = {"stage": _worker_stage, "vcap": _worker_vcap, "nconc": _worker_nconc} + + +def _run_worker_subprocess(mode: str, cfg: dict[str, Any]) -> dict[str, Any]: + """Re-invoke THIS script in a fresh interpreter to run one measurement. + + Fresh process per measurement is load-bearing, not a style choice: ``ru_maxrss`` + only grows within a process, so reusing one process across stages/corpora would + let an earlier, larger measurement's peak silently leak into a later, smaller + one's baseline. + """ + result = subprocess.run( + [sys.executable, str(Path(__file__).resolve()), "--worker", mode, json.dumps(cfg)], + capture_output=True, + text=True, + check=True, + ) + return json.loads(result.stdout.strip().splitlines()[-1]) + + +def _emit(**fields: object) -> None: + print(" ".join(f"{k}={v}" for k, v in fields.items())) + + +# -------------------------------------------------------------------------- +# Orchestration (normal, non-worker invocation) +# -------------------------------------------------------------------------- + + +def _download_corpora(corpora: list[tuple[str, str, str]], cache_dir: Path) -> dict[str, Path]: + cache_dir.mkdir(parents=True, exist_ok=True) + client = httpx.Client(timeout=120.0) + paths: dict[str, Path] = {} + for name, org, repo in corpora: + dest = cache_dir / name + tar_path = dest / "source.tar.gz" + if tar_path.exists(): + print(f"# {name}: using cached {tar_path} ({tar_path.stat().st_size} bytes)") + else: + print(f"# {name}: downloading {org}/{repo}@HEAD ...") + download_tarball(client, org, repo, "HEAD", dest) + print(f"# {name}: downloaded {tar_path} ({tar_path.stat().st_size} bytes)") + paths[name] = tar_path + return paths + + +def main(argv: Sequence[str] | None = None) -> int: + argv = list(sys.argv[1:] if argv is None else argv) + + # Hidden worker-dispatch path: `--worker `. Not a public + # CLI surface -- it exists purely so `_run_worker_subprocess` can re-invoke + # this file in a fresh interpreter for one isolated measurement. + if argv and argv[0] == "--worker": + mode, raw_cfg = argv[1], argv[2] + result = _WORKERS[mode](json.loads(raw_cfg)) + print(json.dumps(result)) + return 0 + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--cache-dir", + type=Path, + default=Path("/tmp/measure_semantic_memory_cache"), + help="where downloaded tarballs are cached across runs", + ) + parser.add_argument( + "--corpora", + default=",".join(f"{org}/{repo}" for _, org, repo in _DEFAULT_CORPORA), + help=( + "comma-separated org/repo list, in size order (last is used for the N-concurrency arm)" + ), + ) + parser.add_argument("--dim", type=int, default=1024, help="embedding dimension") + parser.add_argument("--vcap-n", type=int, default=8000, help="vector count for the V_cap probe") + parser.add_argument( + "--n-threads", default="1,2,3,4", help="comma-separated N values for the concurrency arm" + ) + parser.add_argument( + "--nconc-corpus", + default=None, + help="corpus name to use for the N-concurrency arm (default: the last/largest corpus)", + ) + args = parser.parse_args(argv) + + corpus_specs = [] + for entry in args.corpora.split(","): + org, repo = entry.split("/", 1) + corpus_specs.append((repo, org, repo)) + # Prefer the default's friendly names when they line up (cosmetic only). + default_by_org_repo = {(org, repo): name for name, org, repo in _DEFAULT_CORPORA} + corpus_specs = [ + (default_by_org_repo.get((org, repo), repo), org, repo) for _, org, repo in corpus_specs + ] + + print("=" * 78) + print("environment") + print("=" * 78) + _emit(python=sys.version.split()[0], cpu_count_affinity=_available_cpus()) + try: + import shutil as _shutil + + total, _used, free = _shutil.disk_usage("/tmp") + _emit(tmp_total_bytes=total, tmp_free_bytes=free, note="a-tmpfs-box-per-plan-1.1") + except OSError: + pass + print() + + tarballs = _download_corpora(corpus_specs, args.cache_dir) + print() + + print("=" * 78) + print("stage 1: alpha (files) / gamma (chunks) / vectors, per corpus and gate state") + print("=" * 78) + alphas: list[float] = [] + gammas: list[float] = [] + for name, _org, _repo in corpus_specs: + tarball = tarballs[name] + for gate, delta_fraction in _GATE_STATES: + cfg = {"tarball": str(tarball), "gate": gate, "delta_fraction": delta_fraction} + r = _run_worker_subprocess("stage", cfg) + + files_delta_kb = r["after_files_kb"] - r["baseline_kb"] + chunks_delta_kb = r["after_chunks_kb"] - r["after_files_kb"] + vectors_delta_kb = r["after_vectors_kb"] - r["after_chunks_kb"] + + alpha = (files_delta_kb * 1024) / r["files_bytes"] if r["files_bytes"] else float("nan") + gamma = ( + (chunks_delta_kb * 1024) / r["files_to_embed_bytes"] + if r["files_to_embed_bytes"] + else float("nan") + ) + measured_d = ( + r["chunk_content_bytes"] / r["chunk_count"] if r["chunk_count"] else float("nan") + ) + vector_kb_per_chunk = ( + vectors_delta_kb / r["chunk_count"] if r["chunk_count"] else float("nan") + ) + + _emit( + corpus=name, + gate=gate, + baseline_kb=r["baseline_kb"], + after_files_kb=r["after_files_kb"], + after_chunks_kb=r["after_chunks_kb"], + after_vectors_kb=r["after_vectors_kb"], + files_delta_kb=files_delta_kb, + chunks_delta_kb=chunks_delta_kb, + vectors_delta_kb=vectors_delta_kb, + files_count=r["files_count"], + files_bytes=r["files_bytes"], + files_to_embed_count=r["files_to_embed_count"], + files_to_embed_bytes=r["files_to_embed_bytes"], + chunk_count=r["chunk_count"], + measured_d_bytes_per_chunk=round(measured_d, 1), + vector_kb_per_chunk=round(vector_kb_per_chunk, 3), + alpha_files_per_source_byte=round(alpha, 4), + gamma_chunks_per_source_byte=round(gamma, 4), + ) + + if gate == "first-index": + # alpha/gamma are properties of the whole-file materialization; + # the first-index (gate-closed) run is the one where + # files_to_embed IS files, exactly matching planning's §2.1 + # methodology (a whole-corpus measurement, not a narrowed one). + alphas.append(alpha) + gammas.append(gamma) + print() + + print("=" * 78) + print("stage 2: V_cap -- resident bytes for a fixed embedded-vector count") + print("=" * 78) + vcap_cfg = {"n": args.vcap_n, "dim": args.dim} + vr = _run_worker_subprocess("vcap", vcap_cfg) + vcap_delta_kb = vr["after_kb"] - vr["baseline_kb"] + vcap_kb_per_chunk = vcap_delta_kb / vr["n"] + structural_kb_per_chunk = args.dim * (8 + 24) / 1024 # 8B pointer + 24B PyFloat, per §2.2 + _emit( + n=vr["n"], + dim=vr["dim"], + baseline_kb=vr["baseline_kb"], + after_kb=vr["after_kb"], + delta_kb=vcap_delta_kb, + resident_kb_per_chunk=round(vcap_kb_per_chunk, 3), + structural_kb_per_chunk=round(structural_kb_per_chunk, 3), + ) + print() + + print("=" * 78) + print("stage 3: N concurrent branch-index threads, extraction pool live") + print("=" * 78) + nconc_name = args.nconc_corpus or corpus_specs[-1][0] + nconc_tarball = tarballs[nconc_name] + print(f"# using corpus={nconc_name} tarball={nconc_tarball}, gate=first-index") + n_values = [int(n) for n in args.n_threads.split(",")] + for n in n_values: + cfg = { + "tarball": str(nconc_tarball), + "n_threads": n, + "gate": "first-index", + "delta_fraction": 0.0, + } + nr = _run_worker_subprocess("nconc", cfg) + _emit( + corpus=nconc_name, + n_threads=nr["n_threads"], + n_processes=nr["n_processes"], + baseline_self_kb=nr["baseline_self_kb"], + after_self_kb=nr["after_self_kb"], + self_delta_kb=nr["after_self_kb"] - nr["baseline_self_kb"], + baseline_children_kb=nr["baseline_children_kb"], + after_children_kb=nr["after_children_kb"], + children_delta_kb=nr["after_children_kb"] - nr["baseline_children_kb"], + elapsed_s=round(nr["elapsed_s"], 2), + errors=len(nr["errors"]), + ) + for err in nr["errors"]: + print(f"# error: {err}") + print() + + print("=" * 78) + print("summary") + print("=" * 78) + if alphas: + _emit( + alpha_avg=round(sum(alphas) / len(alphas), 4), + alpha_min=round(min(alphas), 4), + alpha_max=round(max(alphas), 4), + n_corpora=len(alphas), + ) + if gammas: + _emit( + gamma_avg=round(sum(gammas) / len(gammas), 4), + gamma_min=round(min(gammas), 4), + gamma_max=round(max(gammas), 4), + n_corpora=len(gammas), + ) + _emit( + v_cap_resident_kb_per_chunk=round(vcap_kb_per_chunk, 3), + v_cap_structural_kb_per_chunk=round(structural_kb_per_chunk, 3), + v_cap_n=vr["n"], + v_cap_dim=vr["dim"], + ) + print( + "# alpha/gamma above are from each corpus's first-index (gate-closed) run, matching " + "the plan's §2.1 whole-file methodology. The recurring-1pct/10pct rows in stage 1 show " + "the SAME alpha (files always materializes in full) alongside a much smaller " + "chunks/vectors delta (files_to_embed narrows), which is the qualitative effect " + "the plan's §0.2 and §3.3(b) describe." + ) + print( + "# vectors_delta_kb in stage 1 may run slightly high: _precompute_chunk_writer " + "re-chunks internally (it has no seam to accept precomputed chunks), so that delta " + "is 'new allocations since the chunks-term reading' rather than a pure vectors-only " + "measurement. See the module docstring for why this is still a reasonable isolation." + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 8fe93fe..53f7056 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -1187,11 +1187,11 @@ def _embed(texts: list[str]) -> list[list[float]]: def test_config_yaml_semantic_enabled_applies_the_worker_clamp( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: - """The enable overlay also re-arms the 2-worker memory clamp. + """The enable overlay also re-arms the 4-worker memory clamp (issue #109). cfg says disabled (which would leave the pool at index_concurrency=6), but config.yaml's ``semantic.enabled: true`` overlays before effective_workers, so - the clamp fires: pool 6 -> 2, with the clamp log line -- the pool-side mirror + the clamp fires: pool 6 -> 4, with the clamp log line -- the pool-side mirror of the disable test. """ with caplog.at_level(logging.INFO, logger="indexer.job"): @@ -1200,8 +1200,8 @@ def test_config_yaml_semantic_enabled_applies_the_worker_clamp( monkeypatch, cfg=Settings(semantic_enabled=False), ) - assert kwargs["pool_size"] == 2 # clamped - assert "clamping index_concurrency 6 -> 2" in caplog.text + assert kwargs["pool_size"] == 4 # clamped + assert "clamping index_concurrency 6 -> 4" in caplog.text @pytest.mark.unit @@ -2096,10 +2096,10 @@ def test_pool_size_follows_the_semantic_clamp_not_the_raw_config( ) -> None: """The pool must track the EFFECTIVE workers, not index_concurrency. - With semantic on, effective_workers clamps 6 -> 2; a pool of 6 would then - over-provision Lakebase connections that no worker can ever use. The clamp - is also logged, because a run silently doing a third of the requested - concurrency is otherwise invisible. + With semantic on, effective_workers clamps 6 -> 4 (issue #109); a pool of 6 + would then over-provision Lakebase connections that no worker can ever use. + The clamp is also logged, because a run silently doing two-thirds of the + requested concurrency is otherwise invisible. """ with caplog.at_level(logging.INFO, logger="indexer.job"): kwargs = _engine_kwargs( @@ -2107,9 +2107,9 @@ def test_pool_size_follows_the_semantic_clamp_not_the_raw_config( monkeypatch, cfg=Settings(semantic_enabled=True), ) - assert kwargs["pool_size"] == 2 + assert kwargs["pool_size"] == 4 assert kwargs["max_overflow"] == 0 - assert "clamping index_concurrency 6 -> 2" in caplog.text + assert "clamping index_concurrency 6 -> 4" in caplog.text @pytest.mark.unit @@ -2118,7 +2118,7 @@ def test_config_yaml_semantic_disabled_removes_the_worker_clamp( ) -> None: """The overlay runs BEFORE effective_workers, so config.yaml can lift the clamp. - cfg says semantic enabled (which would clamp 6 -> 2), but config.yaml's + cfg says semantic enabled (which would clamp 6 -> 4), but config.yaml's ``semantic.enabled: false`` overlays first -- effective_workers then sees a disabled flag and leaves the pool at the full index_concurrency, with no clamp log line. This is the pool-side proof that the overlay precedes both the clamp @@ -2130,10 +2130,112 @@ def test_config_yaml_semantic_disabled_removes_the_worker_clamp( monkeypatch, cfg=Settings(semantic_enabled=True), ) - assert kwargs["pool_size"] == 6 # not clamped to 2 + assert kwargs["pool_size"] == 6 # not clamped to 4 assert "clamping index_concurrency" not in caplog.text +@pytest.mark.unit +def test_no_clamp_line_at_the_shipped_default( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """At `index_concurrency=4` (the shipped default) the clamp is a no-op. + + `effective_workers` returns `min(index_concurrency, 4)`, and the clamp log + line only fires when it actually reduces the value (`workers != + config.index_concurrency`, job.py). So a default-configured, semantic-on run + emits NO clamp line at all -- issue #109's runbook update explicitly documents + this as confirmed on the live dev job's Arm B runs; this pins it in a unit + test too. + """ + with caplog.at_level(logging.INFO, logger="indexer.job"): + kwargs = _engine_kwargs( + _config(repos=["acme/widgets"], index_concurrency=4), + monkeypatch, + cfg=Settings(semantic_enabled=True), + ) + assert kwargs["pool_size"] == 4 + assert "clamping index_concurrency" not in caplog.text + + +@pytest.mark.unit +def test_shas_fn_connection_closes_before_embedding_starts() -> None: + """E7 (issue #109): the advisory ``shas_fn`` connection is short-lived. + + ``pool_size == workers`` (the tests above) only holds because each worker + opens at most ONE connection at a time -- by sequencing, not by + construction. On the delta-gate-open path a worker opens a SECOND, + short-lived connection (``with engine.connect() as shas_conn:``, + ``indexer/job.py``) for ``shas_fn``, which must close before + ``_precompute_chunk_writer``/embedding starts -- otherwise a single worker + could hold 2 connections at once and pool_size==workers would + under-provision. Wraps the fake engine's own ``connect()`` to count + currently-open connections and samples that count from inside ``shas_fn`` + and ``embed_fn`` -- extending + ``test_shas_fn_called_once_before_embedding_when_version_matches_and_sha_differs``'s + pattern one step further. + """ + open_count = 0 + open_during: dict[str, int] = {} + + engine = _FakeEngine(stamps={("acme/widgets", "main"): ("sha_old", INDEX_SEMANTICS_VERSION)}) + real_connect = engine.connect + + class _TrackingConn: + def __init__(self, inner: _FakeConn) -> None: + self._inner = inner + + def __enter__(self) -> _TrackingConn: + nonlocal open_count + open_count += 1 + self._inner.__enter__() + return self + + def __exit__(self, *exc: Any) -> bool: + nonlocal open_count + open_count -= 1 + return self._inner.__exit__(*exc) + + def execute(self, stmt: Any) -> Any: + return self._inner.execute(stmt) + + def rollback(self) -> None: + self._inner.rollback() + + def _tracking_connect() -> _TrackingConn: + return _TrackingConn(real_connect()) + + engine.connect = _tracking_connect # type: ignore[method-assign] + + def _shas_fn(conn: Any, *, name: str, branch: str) -> tuple[set[Any], set[Any]]: + open_during["shas"] = open_count + return set(), set() + + def _embed(texts: list[str]) -> list[list[float]]: + open_during["embed"] = open_count + return [[0.0] for _ in texts] + + cfg = Settings(semantic_enabled=True, semantic_max_chunks_per_repo=100) + idx = _RecordingIndex() + code = _run( + _config(repos=["acme/widgets"]), + idx, + cfg=cfg, + embed_fn=_embed, + engine=engine, + shas_fn=_shas_fn, + ) + assert code == 0 + # Guards against a vacuous pass: job.py swallows any semantic-precompute + # failure and still returns 0 (the semantic layer is additive), which would + # let a degraded run satisfy the connection-count assertion below without + # ever really reaching the embed step it's supposed to pin. + assert idx.chunk_writer is not None, "semantic precompute must have succeeded, not degraded" + assert open_during == {"shas": 1, "embed": 0}, ( + "shas_fn's own connection must be open exactly while it runs, and fully " + f"closed again before embedding starts (got {open_during!r})" + ) + + @pytest.mark.unit def test_indexer_reaches_no_hardcoded_pool_constant() -> None: """Tripwire: the indexer must never fall back to the server's pool default. diff --git a/tests/unit/test_repo_config.py b/tests/unit/test_repo_config.py index 5861752..d372588 100644 --- a/tests/unit/test_repo_config.py +++ b/tests/unit/test_repo_config.py @@ -253,8 +253,11 @@ def test_extract_processes_out_of_range_raises_config_error(value: int) -> None: ("configured", "semantic_enabled", "expected"), [ (8, False, 8), - (8, True, 2), - (4, True, 2), + (8, True, 4), # issue #109: clamp raised from 2 to 4, re-derived + Arm B-confirmed + (5, True, 4), + (4, True, 4), + (3, True, 3), # below the clamp -- passthrough, N=3 is a legal (unclamped) value too + (2, True, 2), (1, True, 1), # the clamp is a ceiling, never a floor (1, False, 1), ],