bench_run_ledger.py

#!/usr/bin/env python3
"""Measure the 200k-file current-run ledger memory envelope."""

from __future__ import annotations

import argparse
import gc
import json
import logging
import threading
import time
from types import SimpleNamespace

import psutil

from codexa.indexing import changes, scan
from codexa.indexing.routing import _split_queues
from codexa.ingestion.chunking import CHUNKER_VERSION


DEFAULT_FILES = 200_000
DEFAULT_HANDOFF_BUDGET_MIB = 25.0
SYNTHETIC_PATH_BYTES = 100


def _synthetic_path(index: int, extension: str = ".txt") -> str:
    suffix = f"{index:06d}{extension}"
    prefix = "/synthetic/corpus/"
    padding = "x" * (SYNTHETIC_PATH_BYTES - len(prefix) - len(suffix))
    return f"{prefix}{padding}{suffix}"


def _synthetic_meta(count: int, identity_registry):
    for index in range(count):
        if index % 16 == 0:
            identity_registry.claim_dir((2, index // 16))
        if not identity_registry.claim_file((1, index)):
            continue
        yield (
            _synthetic_path(index),
            ".txt",
            SimpleNamespace(
                st_mtime_ns=1_700_000_000_000_000_000 + index,
                st_size=1024 + index % 8192,
            ),
        )


def _manifest(count: int, mode: str) -> dict:
    files = {}
    if mode == "unchanged":
        files = {
            _synthetic_path(index): {
                "hash": f"h{index}",
                "mtime_ns": 1_700_000_000_000_000_000 + index,
                "size": 1024 + index % 8192,
                "chunker_version": CHUNKER_VERSION,
            }
            for index in range(count)
        }
    return {"version": 1, "files": files}


class _PeakRssSampler:
    def __init__(
        self, process: psutil.Process, interval_s: float = 0.005,
    ) -> None:
        self._process = process
        self._interval_s = interval_s
        self._peak = int(process.memory_info().rss)
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, daemon=True)

    def start(self) -> int:
        self._peak = int(self._process.memory_info().rss)
        self._thread.start()
        return self._peak

    def stop(self) -> int:
        self._sample()
        self._stop.set()
        self._thread.join()
        return self._peak

    def _sample(self) -> None:
        self._peak = max(self._peak, int(self._process.memory_info().rss))

    def _run(self) -> None:
        while not self._stop.wait(self._interval_s):
            self._sample()


def measure_run_ledger(count: int, mode: str) -> dict[str, int | float | str]:
    process = psutil.Process()
    process_baseline = int(process.memory_info().rss)
    manifest = _manifest(count, mode)
    gc.collect()
    scan._iter_files_with_meta = (
        lambda *_args, **kwargs: _synthetic_meta(
            count, kwargs["identity_registry"],
        )
    )
    scan._warn_unsupported_exts = lambda *_args, **_kwargs: None
    scan.log_event = lambda *_args, **_kwargs: None
    changes.log_event = lambda *_args, **_kwargs: None
    sampler = _PeakRssSampler(process)
    handoff_baseline = sampler.start()
    started = time.perf_counter()
    inventory = None
    ledger_path = None

    try:
        inventory = scan._scan_phase(
            {"formats": {"enabled": [".txt"]}, "indexer": {}},
            [],
            [],
            logging.getLogger("codexa.mem21_benchmark"),
            return_inventory=True,
        )
        assert isinstance(inventory, scan.ScanInventory)
        new, modified, deleted, hashes = changes.detect_changes(
            inventory.processing_files,
            manifest,
            logging.getLogger("codexa.mem21_benchmark"),
            workers=1,
            physical_presence=inventory.physical_presence,
            sparse_hashes=True,
            disk_backed_changes=True,
        )
        text, ocr = _split_queues(
            inventory.changed_paths,
            total=len(new) + len(modified),
            log=logging.getLogger("codexa.mem21_benchmark"),
        )
        ledger_bytes = inventory.disk_bytes
        ledger_path = inventory.temp_path
    finally:
        peak_rss = sampler.stop()

    expected_changed = count if mode == "fresh" else 0
    if len(new) != expected_changed or modified or deleted or hashes:
        raise RuntimeError("synthetic change classification mismatch")
    if len(text) != expected_changed or ocr:
        raise RuntimeError("synthetic routing mismatch")
    del new, modified, deleted, text, ocr
    inventory.close()
    if ledger_path is not None and ledger_path.exists():
        raise RuntimeError("current-run ledger was not removed")

    mib = 1024 * 1024
    return {
        "files": count,
        "mode": mode,
        "path_bytes": len(_synthetic_path(0).encode()),
        "manifest_mib": round((handoff_baseline - process_baseline) / mib, 2),
        "peak_handoff_mib": round((peak_rss - handoff_baseline) / mib, 2),
        "peak_total_mib": round((peak_rss - process_baseline) / mib, 2),
        "ledger_disk_mib": round(ledger_bytes / mib, 2),
        "duration_s": round(time.perf_counter() - started, 3),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--files", type=int, default=DEFAULT_FILES)
    parser.add_argument(
        "--mode", choices=("fresh", "unchanged"), default="fresh",
    )
    parser.add_argument(
        "--max-handoff-mib",
        type=float,
        default=DEFAULT_HANDOFF_BUDGET_MIB,
    )
    args = parser.parse_args()
    result = measure_run_ledger(args.files, args.mode)
    result["handoff_budget_mib"] = args.max_handoff_mib
    result["within_budget"] = (
        result["peak_handoff_mib"] <= args.max_handoff_mib
    )
    print(json.dumps(result, sort_keys=True))
    return 0 if result["within_budget"] else 1


if __name__ == "__main__":
    raise SystemExit(main())