bench_scan_handoff.py

#!/usr/bin/env python3
"""Measure peak RSS for the synthetic scan/change startup handoff."""

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


DEFAULT_FILES = 200_000
DEFAULT_PEAK_BUDGET_MIB = 80.0
SYNTHETIC_PATH_BYTES = 100


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


def _synthetic_meta(count: int):
    for index in range(count):
        yield (
            _synthetic_path(index),
            ".txt",
            SimpleNamespace(
                st_mtime_ns=1_700_000_000_000_000_000 + index,
                st_size=1024 + index % 8192,
            ),
        )


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._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_scan_handoff(count: int) -> dict[str, int | float]:
    process = psutil.Process()
    scan._iter_files_with_meta = lambda *_args, **_kwargs: _synthetic_meta(count)
    scan._warn_unsupported_exts = lambda *_args, **_kwargs: None
    scan.log_event = lambda *_args, **_kwargs: None
    changes.log_event = lambda *_args, **_kwargs: None
    gc.collect()
    sampler = _PeakRssSampler(process)
    baseline = sampler.start()
    started = time.perf_counter()

    try:
        scanned = scan._scan_phase(
            {"formats": {"enabled": [".txt"]}, "indexer": {}},
            [],
            [],
            logging.getLogger("codexa.scal22_benchmark"),
            return_inventory=True,
        )
        assert isinstance(scanned, scan.ScanInventory)
        new, modified, deleted, hashes = changes.detect_changes(
            scanned.processing_files,
            {"files": {}},
            logging.getLogger("codexa.scal22_benchmark"),
            workers=1,
            physical_presence=scanned.physical_presence,
            sparse_hashes=True,
        )
    finally:
        peak_rss = sampler.stop()

    if len(scanned) != count or len(new) != count:
        raise RuntimeError("synthetic scan dropped files")
    if modified or deleted or hashes:
        raise RuntimeError("fresh synthetic corpus was misclassified")
    ledger_path = scanned.temp_path
    scanned.close()
    if ledger_path.exists():
        raise RuntimeError("scan handoff ledger was not removed")

    peak_mib = max(0, peak_rss - baseline) / (1024 * 1024)
    return {
        "files": count,
        "path_bytes": len(_synthetic_path(0).encode()),
        "peak_incremental_mib": round(peak_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(
        "--max-peak-mib", type=float, default=DEFAULT_PEAK_BUDGET_MIB,
    )
    args = parser.parse_args()
    result = measure_scan_handoff(args.files)
    result["budget_mib"] = args.max_peak_mib
    result["within_budget"] = (
        result["peak_incremental_mib"] <= args.max_peak_mib
    )
    print(json.dumps(result, sort_keys=True))
    return 0 if result["within_budget"] else 1


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