Add explicit source cache populate mode

Author: apstenku123

scripts/streaming_conveyor.py

@@ -1214,6 +1214,38 @@ def stream_lock_names(streams: str) -> tuple[str, ...]:
     return (streams,)
 
 
+def populate_code_source_cache(
+    work_root: Path,
+    source_cache_dir: Path,
+    should_process,
+    progress: ProgressWriter,
+    *,
+    max_repos: int | None = None,
+) -> dict:
+    """Populate the code source cache without running index/tokenize stages."""
+    started = time.monotonic()
+    report = sr.populate_source_cache(
+        work_root,
+        should_process,
+        source_cache_dir,
+        max_repos=max_repos,
+    )
+    for idx, item in enumerate(report["repos"], start=1):
+        progress.emit(
+            "source_cache_repo_ready",
+            repo=item["repo"],
+            repo_dir=item["path"],
+            source_cache_dir=str(source_cache_dir),
+            repo_count=idx,
+        )
+    report = {
+        **report,
+        "elapsed_s": round(time.monotonic() - started, 6),
+    }
+    progress.emit("source_cache_populated", **report)
+    return report
+
+
 # --------------------------------------------------------------------------- #
 # CODE half: orchestrate streaming_reindex.process_one_repo (no reimplementation).
 # Its append_output already lands outputs/reindexed/<L>/<repo>.parquet; we then
@@ -2402,6 +2434,10 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
                    help="For --streams code, only process complete repos already "
                         "in --source-cache-dir; do not open/decompress the source "
                         "tarball.")
+    p.add_argument("--source-cache-populate-only", action="store_true",
+                   help="For --streams code, populate --source-cache-dir from "
+                        "the source tarball and exit without indexing/tokenizing. "
+                        "Follow with --source-cache-only for hot code runs.")
     p.add_argument("--source-dir-root", action="append", default=[],
                    help="For --streams code, process already-extracted repo dirs "
                         "directly without opening the source tarball. May be passed "
@@ -2469,6 +2505,12 @@ def main(argv: list[str]) -> int:
     source_dir_roots = [Path(p) for p in args.source_dir_root]
     if args.source_cache_only and source_cache_dir is None:
         raise SystemExit("--source-cache-only requires --source-cache-dir")
+    if args.source_cache_populate_only and source_cache_dir is None:
+        raise SystemExit("--source-cache-populate-only requires --source-cache-dir")
+    if args.source_cache_populate_only and args.source_cache_only:
+        raise SystemExit("--source-cache-populate-only cannot be combined with --source-cache-only")
+    if args.source_cache_populate_only and args.streams != "code":
+        raise SystemExit("--source-cache-populate-only is only valid with --streams code")
     if args.source_cache_only and args.streams != "code":
         raise SystemExit("--source-cache-only is only valid with --streams code")
     if source_cache_dir is not None and args.streams != "code":
@@ -2659,6 +2701,7 @@ def _on_signal(signum, _frame):
         code_recompress_workers=args.code_recompress_workers if code_recompressor else 0,
         source_cache_dir=str(source_cache_dir) if source_cache_dir is not None else None,
         source_cache_only=args.source_cache_only,
+        source_cache_populate_only=args.source_cache_populate_only,
         source_dir_roots=[str(p) for p in source_dir_roots],
         reservation_file=str(args.reservation_file)
         if reservations is not None else None,
@@ -2737,6 +2780,42 @@ def should_process(repo: str) -> bool:
             )
         return should_stage
 
+    if args.source_cache_populate_only:
+        try:
+            report = populate_code_source_cache(
+                work_root,
+                source_cache_dir,
+                should_process,
+                progress,
+                max_repos=args.max_repos,
+            )
+            summary = {
+                "repos_this_run": 0,
+                "code_halves_this_run": 0,
+                "commit_ranges_this_run": 0,
+                "commit_ranges_failed_this_run": 0,
+                "workers": workers,
+                "repo_workers": repo_workers,
+                "max_active_repos": max_active_repos,
+                "parse_workers": parse_workers,
+                "memory_plan": memory_plan,
+                "streams": args.streams,
+                "source_cache_populate_only": True,
+                "source_cache_report": report,
+                "manifest": str(CONVEYOR_MANIFEST),
+                "interrupted": STOP_EVENT.is_set(),
+            }
+            progress.emit("run_finished", **summary)
+            print(json.dumps(summary, indent=2))
+            return 130 if STOP_EVENT.is_set() else 0
+        finally:
+            if code_recompressor is not None:
+                code_recompressor.shutdown()
+            if own_work_root and not args.keep_temp:
+                remove_tree(work_root, reason="source-cache populate work_root")
+            for lock in reversed(run_locks):
+                lock.close()
+
     range_pool = ThreadPoolExecutor(max_workers=workers)
     repo_pool = ThreadPoolExecutor(max_workers=repo_workers) if repo_workers > 1 else None
 

scripts/streaming_reindex.py

@@ -659,6 +659,43 @@ def stream_repo_dirs(source_roots: Sequence[Path], should_process):
             yield repo, repo_dir
 
 
+def populate_source_cache(
+    work_root: Path,
+    should_process,
+    source_cache_dir: Path,
+    *,
+    max_repos: int | None = None,
+) -> dict:
+    """Materialize source-cache repos without running the tokenization pipeline.
+
+    This is the explicit "cold source-store build" phase. It may read the
+    monolithic source tar once, but callers can then run the hot code conveyor
+    with --source-cache-only and avoid any sequential tar drain in production.
+    """
+    if source_cache_dir is None:
+        raise ValueError("source_cache_dir is required")
+    repos: list[dict[str, str]] = []
+    gen = stream_repo_subtrees(
+        work_root,
+        should_process,
+        source_cache_dir=source_cache_dir,
+        source_cache_only=False,
+    )
+    try:
+        for repo, repo_dir in gen:
+            repos.append({"repo": repo, "path": str(repo_dir)})
+            if max_repos is not None and len(repos) >= max_repos:
+                break
+    finally:
+        if hasattr(gen, "close"):
+            gen.close()
+    return {
+        "source_cache_dir": str(source_cache_dir),
+        "repos": repos,
+        "repo_count": len(repos),
+    }
+
+
 # --------------------------------------------------------------------------- #
 # Per-repo pipeline stages.
 # --------------------------------------------------------------------------- #
@@ -1070,6 +1107,10 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
     p.add_argument("--source-cache-only", action="store_true",
                    help="Only process repos already complete in --source-cache-dir; "
                         "do not open/decompress the source tarball.")
+    p.add_argument("--source-cache-populate-only", action="store_true",
+                   help="Populate --source-cache-dir from the source tarball and "
+                        "exit without indexing/tokenizing. Use the resulting cache "
+                        "with --source-cache-only for hot code runs.")
     p.add_argument("--source-dir-root", action="append", default=[],
                    help="Already-extracted repo root to process directly, without "
                         "opening the source tarball. May be passed multiple times; "
@@ -1186,8 +1227,25 @@ def should_process(repo: str) -> bool:
         source_dir_roots = [Path(p) for p in args.source_dir_root]
         if args.source_cache_only and source_cache_dir is None:
             raise SystemExit("--source-cache-only requires --source-cache-dir")
+        if args.source_cache_populate_only and source_cache_dir is None:
+            raise SystemExit("--source-cache-populate-only requires --source-cache-dir")
+        if args.source_cache_populate_only and args.source_cache_only:
+            raise SystemExit("--source-cache-populate-only cannot be combined with --source-cache-only")
+        if args.source_cache_populate_only and args.commit_source:
+            raise SystemExit("--source-cache-populate-only is code-only; remove --commit-source")
         if source_dir_roots and (source_cache_dir is not None or args.source_cache_only):
             raise SystemExit("--source-dir-root cannot be combined with source cache flags")
+        if args.source_cache_populate_only:
+            report = populate_source_cache(
+                work_root,
+                should_process,
+                source_cache_dir,
+                max_repos=args.max_repos,
+            )
+            print(json.dumps(report, indent=2, sort_keys=True))
+            if own_work_root and not args.keep_temp:
+                shutil.rmtree(work_root, ignore_errors=True)
+            return 0
         gen = (
             stream_repo_dirs(source_dir_roots, should_process)
             if source_dir_roots

tests/test_streaming_conveyor_progress.py

@@ -58,6 +58,59 @@ def test_progress_writer_tracks_extract_cache_rates(tmp_path):
     assert writer.extract_cache_metrics()["status_counts"]["hit"] == 1
 
 
+def test_populate_code_source_cache_emits_ready_and_summary(tmp_path, monkeypatch):
+    import streaming_conveyor
+
+    progress_path = tmp_path / "progress.jsonl"
+    progress = streaming_conveyor.ProgressWriter(progress_path)
+    cache = tmp_path / "source_cache"
+    observed = {}
+
+    def fake_populate(work_root, should_process, source_cache_dir, *, max_repos):
+        observed["work_root"] = work_root
+        observed["source_cache_dir"] = source_cache_dir
+        observed["max_repos"] = max_repos
+        repos = []
+        for repo in ("repo-a", "repo.bare", "repo-b"):
+            if should_process(repo):
+                repos.append({"repo": repo, "path": str(source_cache_dir / repo)})
+        repos = repos[:max_repos]
+        return {
+            "source_cache_dir": str(source_cache_dir),
+            "repos": repos,
+            "repo_count": len(repos),
+        }
+
+    monkeypatch.setattr(streaming_conveyor.sr, "populate_source_cache", fake_populate)
+
+    report = streaming_conveyor.populate_code_source_cache(
+        tmp_path / "work",
+        cache,
+        streaming_conveyor.sr.is_code_worktree_repo,
+        progress,
+        max_repos=1,
+    )
+
+    assert observed == {
+        "work_root": tmp_path / "work",
+        "source_cache_dir": cache,
+        "max_repos": 1,
+    }
+    assert report["repo_count"] == 1
+    rows = [
+        json.loads(line)
+        for line in progress_path.read_text(encoding="utf-8").splitlines()
+    ]
+    assert [row["event"] for row in rows] == [
+        "source_cache_repo_ready",
+        "source_cache_populated",
+    ]
+    assert rows[0]["repo"] == "repo-a"
+    assert rows[0]["repo_dir"] == str(cache / "repo-a")
+    assert rows[1]["repo_count"] == 1
+    assert rows[1]["repos"] == [{"repo": "repo-a", "path": str(cache / "repo-a")}]
+
+
 def test_background_recompressor_runs_and_surfaces_completion(tmp_path, monkeypatch):
     import streaming_conveyor
 

tests/test_streaming_reindex_run_checked.py

@@ -154,3 +154,32 @@ def test_stream_repo_dirs_yields_extracted_src_without_tar(tmp_path) -> None:
     )
 
     assert yielded == [("repo", root / "repo" / "_src")]
+
+
+def test_populate_source_cache_only_materializes_cache(tmp_path, monkeypatch) -> None:
+    import streaming_reindex
+
+    cache = tmp_path / "cache"
+    calls = []
+
+    def fake_stream(work_root, should_process, *, source_cache_dir, source_cache_only):
+        calls.append((work_root, source_cache_dir, source_cache_only))
+        for repo in ("repo-a", "repo.bare", "repo-b"):
+            if should_process(repo):
+                yield repo, source_cache_dir / repo
+
+    monkeypatch.setattr(streaming_reindex, "stream_repo_subtrees", fake_stream)
+
+    report = streaming_reindex.populate_source_cache(
+        tmp_path / "work",
+        streaming_reindex.is_code_worktree_repo,
+        cache,
+        max_repos=1,
+    )
+
+    assert calls == [(tmp_path / "work", cache, False)]
+    assert report == {
+        "source_cache_dir": str(cache),
+        "repos": [{"repo": "repo-a", "path": str(cache / "repo-a")}],
+        "repo_count": 1,
+    }