cleveragents-core/benchmarks/large_project_scaling_bench.py

"""ASV benchmarks for large-project indexing at production scale.

Measures the performance of ``walk_and_index`` and incremental refresh
at 1K, 10K, 50K, and 100K file counts, tracking throughput (files and
tokens per second) and indexed-file totals.

Extends the existing 5K-file ceiling from
:mod:`large_project_decompose_bench` to production-scale repositories.
"""

from __future__ import annotations

import importlib
import os
import random
import shutil
import sys
import tempfile
import time
from pathlib import Path
from typing import ClassVar

# Ensure the local *source* tree is importable even when ASV has an
# older build of the package installed.
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

import cleveragents  # noqa: E402

importlib.reload(cleveragents)

from cleveragents.application.services.repo_indexing_utils import (  # noqa: E402
    walk_and_index,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

# Realistic directory structure skeleton used for synthetic repos.
_DIRS = ("src", "tests", "docs", "scripts", "config")

# Extension → content template.  Keeps files realistic so that
# language detection and token estimation exercise real code paths.
_EXTENSIONS: dict[str, str] = {
    ".py": "# auto-generated\ndef func_{i}() -> int:\n    return {i}\n",
    ".md": "# Document {i}\n\nGenerated documentation paragraph.\n",
    ".json": '{{"id": {i}, "name": "item_{i}"}}\n',
    ".ts": "export const value_{i}: number = {i};\n",
    ".yaml": "key_{i}: value_{i}\n",
}

_EXT_LIST = list(_EXTENSIONS.keys())


def _build_project(root: str, file_count: int) -> None:
    """Populate *root* with *file_count* files in a realistic layout."""
    rng = random.Random(42)  # deterministic seed
    for i in range(file_count):
        subdir = _DIRS[i % len(_DIRS)]
        # Create two levels of nesting to mimic real projects.
        nested = f"pkg{i % 20}"
        dirpath = os.path.join(root, subdir, nested)
        os.makedirs(dirpath, exist_ok=True)
        ext = _EXT_LIST[rng.randint(0, len(_EXT_LIST) - 1)]
        template = _EXTENSIONS[ext]
        content = template.format(i=i)
        fpath = os.path.join(dirpath, f"file_{i:06d}{ext}")
        with open(fpath, "w") as fh:
            fh.write(content)


def _modify_subset(root: str, file_count: int, pct: float = 0.01) -> int:
    """Modify *pct* of files under *root* and return count modified."""
    rng = random.Random(99)  # deterministic seed
    modify_count = max(1, int(file_count * pct))
    indices = rng.sample(range(file_count), modify_count)
    modified = 0
    for i in indices:
        subdir = _DIRS[i % len(_DIRS)]
        nested = f"pkg{i % 20}"
        ext = _EXT_LIST[rng.randint(0, len(_EXT_LIST) - 1)]
        fpath = os.path.join(root, subdir, nested, f"file_{i:06d}{ext}")
        if os.path.exists(fpath):
            with open(fpath, "a") as fh:
                fh.write(f"\n# modified at {time.monotonic()}\n")
            modified += 1
    return modified


# ---------------------------------------------------------------------------
# Parameterized indexing suite
# ---------------------------------------------------------------------------

_DEFAULT_INCLUDE: tuple[str, ...] = ()
_DEFAULT_EXCLUDE: tuple[str, ...] = ("*.pyc", "__pycache__/*")


class IndexingScalingSuite:
    """Benchmark ``walk_and_index`` at production-scale file counts."""

    params: ClassVar[list[int]] = [1_000, 10_000, 50_000, 100_000]
    param_names: ClassVar[list[str]] = ["file_count"]
    timeout = 600  # 10 min for 100K files

    _tmpdir: str

    def setup(self, file_count: int) -> None:
        """Create a temp directory with *file_count* files."""
        self._tmpdir = tempfile.mkdtemp(prefix="bench-scale-idx-")
        _build_project(self._tmpdir, file_count)

    def teardown(self, file_count: int) -> None:
        shutil.rmtree(self._tmpdir, ignore_errors=True)

    # -- timing methods -----------------------------------------------------

    def time_walk_and_index(self, file_count: int) -> None:
        """Time full ``walk_and_index`` for *file_count* files."""
        walk_and_index(
            root=Path(self._tmpdir),
            include_globs=_DEFAULT_INCLUDE,
            exclude_globs=_DEFAULT_EXCLUDE,
            max_file_size=None,
            max_total_size=None,
        )

    def time_incremental_refresh(self, file_count: int) -> None:
        """Modify 1 %% of files, then re-index the full tree.

        This measures incremental *walk* overhead; the actual diff merge
        happens in ``RepoIndexingService.refresh_index`` which requires
        a database.  Here we just re-walk to quantify I/O cost.
        """
        _modify_subset(self._tmpdir, file_count, pct=0.01)
        walk_and_index(
            root=Path(self._tmpdir),
            include_globs=_DEFAULT_INCLUDE,
            exclude_globs=_DEFAULT_EXCLUDE,
            max_file_size=None,
            max_total_size=None,
        )

    # -- tracking methods ---------------------------------------------------

    def track_indexed_file_count(self, file_count: int) -> int:
        """Track the number of files actually indexed."""
        records = walk_and_index(
            root=Path(self._tmpdir),
            include_globs=_DEFAULT_INCLUDE,
            exclude_globs=_DEFAULT_EXCLUDE,
            max_file_size=None,
            max_total_size=None,
        )
        return len(records)

    def track_tokens_per_second(self, file_count: int) -> float:
        """Track indexing throughput in tokens / second."""
        t0 = time.perf_counter()
        records = walk_and_index(
            root=Path(self._tmpdir),
            include_globs=_DEFAULT_INCLUDE,
            exclude_globs=_DEFAULT_EXCLUDE,
            max_file_size=None,
            max_total_size=None,
        )
        elapsed = time.perf_counter() - t0
        total_tokens = sum(r.token_count for r in records)
        if elapsed <= 0:
            return float(total_tokens)
        return total_tokens / elapsed


# Attach ASV unit metadata without ``# type: ignore``.
setattr(  # noqa: B010
    IndexingScalingSuite.track_indexed_file_count, "unit", "files"
)
setattr(  # noqa: B010
    IndexingScalingSuite.track_tokens_per_second, "unit", "tokens/s"
)