DocForge2/tools/milestone1_benchmark.py

"""Milestone 1 warm-operation benchmark with algorithmic zero-work gates."""

from __future__ import annotations

import argparse
import json
import platform
import resource
import subprocess
import sys
import tempfile
import threading
import time
from collections.abc import Callable, Mapping
from pathlib import Path
from typing import cast

from milestone0_baseline import (
    measure_operation,
    synthetic_node_id,
    write_synthetic_project,
)

from docforge.context import compile_context
from docforge.index import ProjectIndex
from docforge.mcp_server import DocForgeService
from docforge.project import Project
from docforge.rendering import RenderService
from docforge.viewer_manager import ViewerManager, ViewerManagerClient

ROOT = Path(__file__).resolve().parents[1]
ZERO_WORK_COUNTERS = (
    "project_loads",
    "source_files_parsed",
    "source_bytes_parsed",
    "adapter_projection_loads",
    "adapter_source_extractions",
    "index_builds",
    "render_prepare_calls",
    "render_output_bytes_built",
    "render_output_bytes_hashed",
)


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Gate warm DocForge2 core work on a disposable deterministic project."
    )
    parser.add_argument("--nodes", type=int, default=1000)
    parser.add_argument("--samples", type=int, default=10)
    parser.add_argument("--output", type=Path)
    return parser


def _git(command: list[str]) -> str:
    return subprocess.run(
        ["git", *command],
        cwd=ROOT,
        check=True,
        capture_output=True,
        text=True,
    ).stdout.strip()


def _diagnostics(result: object) -> Mapping[str, object]:
    if not isinstance(result, Mapping):
        raise RuntimeError("Measured operation returned a non-object result")
    result_payload = cast(Mapping[str, object], result)
    diagnostics_value = result_payload.get("diagnostics")
    if not isinstance(diagnostics_value, Mapping):
        raise RuntimeError("Measured operation did not return diagnostics")
    diagnostics = cast(Mapping[str, object], diagnostics_value)
    counters_value = diagnostics.get("counters")
    if not isinstance(counters_value, Mapping):
        raise RuntimeError("Measured diagnostics did not return counters")
    counters = cast(Mapping[str, object], counters_value)
    for counter in ZERO_WORK_COUNTERS:
        if counters.get(counter) != 0:
            raise RuntimeError(f"Warm operation performed forbidden work: {counter}")
    return diagnostics


def _operation(
    operation: Callable[[], dict[str, object]],
    *,
    samples: int,
    p95_limit_ms: float,
    expected_status: str = "ok",
) -> dict[str, object]:
    measurement, last = measure_operation(operation, samples=samples)
    if not isinstance(last, Mapping):
        raise RuntimeError(f"Measured operation did not return status={expected_status}")
    last_payload = cast(Mapping[str, object], last)
    if last_payload.get("status") != expected_status:
        raise RuntimeError(f"Measured operation did not return status={expected_status}")
    diagnostics = _diagnostics(last_payload)
    p95_ms = float(cast(float, measurement["p95_ms"]))
    if p95_ms > p95_limit_ms:
        raise RuntimeError(f"Warm operation p95 {p95_ms:.3f} ms exceeds {p95_limit_ms:.3f} ms")
    return {
        **measurement,
        "p95_limit_ms": p95_limit_ms,
        "diagnostics": diagnostics,
    }


def _benchmark(root: Path, node_count: int, samples: int) -> dict[str, object]:
    project = Project.open(root)
    ProjectIndex(project).build()
    RenderService(project).render("manual")
    service = DocForgeService(project, diagnostics=True)
    target = synthetic_node_id(node_count - 1)
    operations = {
        "warm_no_change_synchronize": _operation(
            service.synchronize,
            samples=samples,
            p95_limit_ms=100,
        ),
        "exact_node": _operation(
            lambda: service.invoke(
                lambda: service.index.get_node(target),
                operation_name="mcp.get_node",
            ),
            samples=samples,
            p95_limit_ms=50,
        ),
        "missing_node_error": _operation(
            lambda: service.invoke(
                lambda: service.index.get_node("missing.node"),
                operation_name="mcp.get_node",
            ),
            samples=samples,
            p95_limit_ms=50,
            expected_status="error",
        ),
        "search_limit_20": _operation(
            lambda: service.invoke(
                lambda: service.index.search("Synthetic measurement", limit=20),
                operation_name="mcp.search",
            ),
            samples=samples,
            p95_limit_ms=100,
        ),
        "dependencies_depth_8": _operation(
            lambda: service.invoke(
                lambda: service.index.dependencies(target, depth=8, limit=100),
                operation_name="mcp.dependencies",
            ),
            samples=samples,
            p95_limit_ms=100,
        ),
        "context_32k": _operation(
            lambda: service.invoke(
                lambda: compile_context(service.index, "active", 32_000),
                operation_name="mcp.context",
            ),
            samples=samples,
            p95_limit_ms=250,
        ),
        "render_receipt_status": _operation(
            lambda: service.render_status("manual"),
            samples=samples,
            p95_limit_ms=50,
        ),
    }
    state_path = root / ".docforge" / "benchmark-viewer-manager.json"
    manager = ViewerManager(state_path, check_interval_seconds=0.02)
    manager_thread = threading.Thread(target=manager.serve_forever, daemon=True)
    manager_thread.start()
    deadline = time.monotonic() + 2
    while not state_path.exists() and time.monotonic() < deadline:
        time.sleep(0.01)
    if not state_path.exists():
        manager.shutdown()
        manager_thread.join(timeout=2)
        raise RuntimeError("Viewer manager did not start")
    service.visualization = ViewerManagerClient(service.index, state_path=state_path)
    try:
        service.visualization.start()
        operations["visualization_current_status"] = _operation(
            service.visualization_status,
            samples=samples,
            p95_limit_ms=50,
        )
        with service.index.path.open("ab") as stream:
            stream.write(b"\n")
        operations["visualization_stale_status"] = _operation(
            service.visualization_status,
            samples=samples,
            p95_limit_ms=50,
        )
        service.stop_visualization()
        operations["visualization_not_running_status"] = _operation(
            service.visualization_status,
            samples=samples,
            p95_limit_ms=50,
        )
    finally:
        manager.shutdown()
        manager_thread.join(timeout=2)
    service.visualization = ViewerManagerClient(service.index, state_path=state_path)
    operations["visualization_unavailable_status"] = _operation(
        service.visualization_status,
        samples=samples,
        p95_limit_ms=50,
        expected_status="error",
    )
    return {
        "fixture": {
            "kind": "synthetic_generic",
            "node_count": node_count,
            "edge_count": node_count - 1,
            "source_file_count": node_count,
        },
        "operations": operations,
        "process_peak_rss_kib": int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss),
    }


def main() -> int:
    arguments = _parser().parse_args()
    if arguments.nodes < 2:
        raise SystemExit("--nodes must be at least 2")
    if arguments.samples < 1:
        raise SystemExit("--samples must be positive")
    with tempfile.TemporaryDirectory(prefix="docforge-milestone1-") as directory:
        root = Path(directory).resolve()
        write_synthetic_project(root, arguments.nodes)
        measurement = _benchmark(root, arguments.nodes, arguments.samples)
    result = {
        "schema_version": 1,
        "benchmark": "docforge2_milestone1",
        "source": {
            "revision": _git(["rev-parse", "HEAD"]),
            "dirty": bool(_git(["status", "--porcelain"])),
        },
        "environment": {
            "platform": platform.platform(),
            "machine": platform.machine(),
            "python": platform.python_version(),
            "implementation": platform.python_implementation(),
        },
        "method": {
            "clock": "time.perf_counter_ns",
            "memory": "resource.getrusage(RUSAGE_SELF).ru_maxrss",
            "response_size": "UTF-8 bytes of compact sorted JSON",
            "samples": arguments.samples,
            "zero_work_counters": list(ZERO_WORK_COUNTERS),
        },
        **measurement,
    }
    encoded = json.dumps(result, sort_keys=True, indent=2) + "\n"
    if arguments.output is not None:
        output = arguments.output.resolve()
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(encoded, encoding="utf-8")
    sys.stdout.write(encoded)
    return 0


if __name__ == "__main__":
    sys.exit(main())