from __future__ import annotations
import argparse
import json
import os
import platform
import signal
import shutil
import statistics
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import TextIO
ROOT = Path(__file__).resolve().parents[2]
CONFIG = Path(__file__).with_name("bot-forge.toml")
DEFAULT_BINARY = ROOT / "target" / "release" / "bot-forge"
DEFAULT_OUTPUT = ROOT / "target" / "benchmarks" / "quality-install"
TOOLCHAIN = "1.89.0"
RSPROXY_INDEX = "sparse+https://rsproxy.cn/index/"
RSPROXY_SOURCE = "rsproxy-sparse"
CARGO_TOOLS = (
("cargo-expand", "1.0.126", None, None),
("cargo-nextest", "0.9.120", None, None),
("cargo-audit", "0.22.2", None, None),
("cargo-deny", "0.20.2", None, None),
("cargo-geiger", "0.13.0", None, None),
("cargo-llvm-cov", "0.9.0", None, None),
(
"bot-gate",
"1.2.1",
None,
None,
),
("bot-metric", "1.2.0", None, None),
)
VERIFY = (
("rustfmt", ["rustfmt", f"+{TOOLCHAIN}", "--version"]),
("clippy", ["cargo", f"+{TOOLCHAIN}", "clippy", "--version"]),
(
"llvm-tools-preview",
["rustup", "component", "list", "--installed", "--toolchain", TOOLCHAIN],
),
("cargo-expand", ["cargo", f"+{TOOLCHAIN}", "expand", "--version"]),
("cargo-nextest", ["cargo", f"+{TOOLCHAIN}", "nextest", "--version"]),
("cargo-audit", ["cargo", f"+{TOOLCHAIN}", "audit", "--version"]),
("cargo-deny", ["cargo", f"+{TOOLCHAIN}", "deny", "--version"]),
("cargo-geiger", ["cargo", f"+{TOOLCHAIN}", "geiger", "--version"]),
("cargo-llvm-cov", ["cargo", f"+{TOOLCHAIN}", "llvm-cov", "--version"]),
("bot-gate", ["bot-gate", "--version"]),
("bot-metric", ["bot-metric", "--version"]),
)
@dataclass
class Sample:
method: str
round: int
order: int
duration_seconds: float
peak_rss_mib: float
log: str
versions: dict[str, str]
def checked(command: list[str], env: dict[str, str] | None = None) -> str:
result = subprocess.run(
command,
cwd=ROOT,
env=env,
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
check=False,
)
if result.returncode:
raise RuntimeError(f"command failed ({result.returncode}): {' '.join(command)}\n{result.stdout}")
return result.stdout.strip()
def process_tree_rss_mib(root_pid: int) -> float:
result = subprocess.run(
["ps", "-axo", "pid=,ppid=,rss="], capture_output=True, text=True, check=False
)
rows = []
for line in result.stdout.splitlines():
fields = line.split()
if len(fields) == 3:
rows.append(tuple(map(int, fields)))
descendants = {root_pid}
changed = True
while changed:
changed = False
for pid, parent, _ in rows:
if parent in descendants and pid not in descendants:
descendants.add(pid)
changed = True
return sum(rss for pid, _, rss in rows if pid in descendants) / 1024
def run_measured(
command: list[str], env: dict[str, str], log: TextIO, cwd: Path
) -> tuple[float, float]:
started = time.perf_counter()
process = subprocess.Popen(
command,
cwd=cwd,
env=env,
stdin=subprocess.DEVNULL,
stdout=log,
stderr=subprocess.STDOUT,
text=True,
start_new_session=True,
)
peak = 0.0
try:
while process.poll() is None:
peak = max(peak, process_tree_rss_mib(process.pid))
time.sleep(0.25)
except BaseException:
if process.poll() is None:
os.killpg(process.pid, signal.SIGTERM)
try:
process.wait(timeout=5)
except subprocess.TimeoutExpired:
os.killpg(process.pid, signal.SIGKILL)
process.wait()
raise
elapsed = time.perf_counter() - started
if process.returncode:
raise RuntimeError(f"command failed ({process.returncode}): {' '.join(command)}")
return elapsed, peak
def filtered_system_path(rustup_bin: Path) -> str:
rejected = (
"/.cargo/bin",
"/bot-forge/bin",
"/Application Support/bot-forge/bin",
"/opt/homebrew/opt/rustup/bin",
)
entries = [
entry
for entry in os.environ.get("PATH", "").split(os.pathsep)
if entry and not any(marker in entry for marker in rejected)
]
return os.pathsep.join([str(rustup_bin), *entries])
def sample_env(
sample_root: Path,
rustup_bin: Path,
rustup_home: Path,
method: str,
cpu_count: int,
) -> dict[str, str]:
home = sample_root / "home"
cargo_home = sample_root / "cargo-home"
forge_home = sample_root / "forge-home"
serial_root = sample_root / "serial-root"
for path in [home, cargo_home, forge_home, serial_root]:
path.mkdir(parents=True, exist_ok=True)
cargo_config = sample_root / ".cargo" / "config.toml"
cargo_config.parent.mkdir(parents=True, exist_ok=True)
cargo_config.write_text(
'[source.crates-io]\n'
f'replace-with = "{RSPROXY_SOURCE}"\n\n'
f'[source.{RSPROXY_SOURCE}]\n'
f'registry = "{RSPROXY_INDEX}"\n\n'
'[net]\n'
'git-fetch-with-cli = true\n',
encoding="utf-8",
)
first_bin = forge_home / "bin" if method == "bot-forge" else serial_root / "bin"
path = os.pathsep.join([str(first_bin), str(cargo_home / "bin"), filtered_system_path(rustup_bin)])
env = os.environ.copy()
env.update(
{
"HOME": str(home),
"USERPROFILE": str(home),
"RUSTUP_HOME": str(rustup_home),
"CARGO_HOME": str(cargo_home),
"BOT_FORGE_HOME": str(forge_home),
"BOT_FORGE_BIN_DIR": str(forge_home / "bin"),
"BOT_FORGE_CONFIG_DIR": str(sample_root / "forge-config"),
"CARGO_BUILD_JOBS": str(cpu_count),
"CARGO_TERM_COLOR": "never",
"NO_COLOR": "1",
"PATH": path,
}
)
return env
def verify_rust_components(env: dict[str, str]) -> dict[str, str]:
versions = {}
for name, command in VERIFY[:3]:
output = checked(command, env)
if name == "llvm-tools-preview":
matches = [line for line in output.splitlines() if line.startswith("llvm-tools")]
if not matches:
raise RuntimeError("llvm-tools-preview is not installed in Rust 1.89")
output = matches[0]
versions[name] = output.splitlines()[0]
return versions
def verify_rsproxy_cache(sample_root: Path, method: str) -> None:
if method == "bot-forge":
cache_root = sample_root / "forge-home" / "cache" / "cargo" / "sources"
else:
cache_root = sample_root / "cargo-home"
configs = list(cache_root.rglob("registry/index/*/config.json"))
if not configs:
raise RuntimeError(f"{method} did not create a Cargo registry index cache")
download_urls = {
str(json.loads(config.read_text(encoding="utf-8")).get("dl", ""))
for config in configs
}
if not download_urls or any(
not url.startswith("https://rsproxy.cn/") for url in download_urls
):
raise RuntimeError(
f"{method} used a non-rsproxy Cargo download endpoint: {sorted(download_urls)}"
)
def serial_install(sample_root: Path, env: dict[str, str], log: TextIO) -> tuple[float, float]:
commands = [command for _, command in VERIFY[:3]]
root = sample_root / "serial-root"
for name, version, source, revision in CARGO_TOOLS:
command = [
"cargo",
f"+{TOOLCHAIN}",
"install",
name,
"--version",
f"={version}",
"--locked",
"--root",
str(root),
]
if source:
command = [
"cargo",
f"+{TOOLCHAIN}",
"install",
name,
"--version",
f"={version}",
"--git",
source,
"--rev",
revision or "",
"--locked",
"--root",
str(root),
]
commands.append(command)
started = time.perf_counter()
peak = 0.0
for command in commands:
_, command_peak = run_measured(command, env, log, sample_root)
peak = max(peak, command_peak)
return time.perf_counter() - started, peak
def verify(env: dict[str, str]) -> dict[str, str]:
versions = {}
for name, command in VERIFY:
output = checked(command, env)
if name == "llvm-tools-preview":
matches = [line for line in output.splitlines() if line.startswith("llvm-tools")]
if not matches:
raise RuntimeError("llvm-tools-preview verification failed")
output = matches[0]
versions[name] = output.splitlines()[0]
return versions
def run_sample(
method: str,
round_number: int,
order: int,
workspace: Path,
binary: Path,
rustup_bin: Path,
rustup_home: Path,
cpu_count: int,
log_dir: Path,
) -> Sample:
sample_root = workspace / f"round-{round_number:02d}-{order}-{method}"
sample_root.mkdir(parents=True)
env = sample_env(sample_root, rustup_bin, rustup_home, method, cpu_count)
if any((sample_root / "cargo-home").iterdir()) or any(
(sample_root / "forge-home").iterdir()
):
raise RuntimeError(f"{method} sample did not start with empty caches")
log_path = log_dir / f"round-{round_number:02d}-{order}-{method}.log"
with log_path.open("w", encoding="utf-8") as log:
if method == "bot-forge":
duration, peak = run_measured(
[str(binary), "install", "bench-quality", "--yes", "--quiet", "--config", str(CONFIG)],
env,
log,
sample_root,
)
else:
duration, peak = serial_install(sample_root, env, log)
verify_rsproxy_cache(sample_root, method)
versions = verify(env)
return Sample(method, round_number, order, duration, peak, str(log_path), versions)
def summary(samples: list[Sample], method: str) -> dict[str, float | int]:
values = [sample.duration_seconds for sample in samples if sample.method == method]
peaks = [sample.peak_rss_mib for sample in samples if sample.method == method]
return {
"samples": len(values),
"min_seconds": round(min(values), 3),
"median_seconds": round(statistics.median(values), 3),
"max_seconds": round(max(values), 3),
"median_peak_rss_mib": round(statistics.median(peaks), 1),
}
def markdown(report: dict[str, object]) -> str:
methods = report["summary"]
forge = methods["bot-forge"]
serial = methods["serial"]
lines = [
"# Quality installation benchmark result",
"",
f"- Host: {report['host']['platform']} / {report['host']['machine']} / {report['host']['logical_cpus']} logical CPUs",
f"- Cargo registry: {report['cargo_registry']}",
"- Cargo cache: empty and isolated for every sample",
"- Rust 1.89 components: current-machine detection, included in end-to-end timing",
f"- Rounds: {report['runs']}",
"",
"| Method | Median | Min | Max | Median peak RSS |",
"| --- | ---: | ---: | ---: | ---: |",
f"| bot-forge | {forge['median_seconds']:.3f} s | {forge['min_seconds']:.3f} s | {forge['max_seconds']:.3f} s | {forge['median_peak_rss_mib']:.1f} MiB |",
f"| serial | {serial['median_seconds']:.3f} s | {serial['min_seconds']:.3f} s | {serial['max_seconds']:.3f} s | {serial['median_peak_rss_mib']:.1f} MiB |",
"",
f"**Speedup: {report['speedup']:.3f}×** (`serial median / bot-forge median`)",
"",
"| Round | Order | Method | Wall time | Peak RSS |",
"| ---: | ---: | --- | ---: | ---: |",
]
for sample in report["samples"]:
lines.append(
f"| {sample['round']} | {sample['order']} | {sample['method']} | "
f"{sample['duration_seconds']:.3f} s | {sample['peak_rss_mib']:.1f} MiB |"
)
lines.extend(["", "All 11 requested targets were verified after every successful sample.", ""])
return "\n".join(lines)
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--runs", type=int, default=1)
parser.add_argument("--binary", type=Path, default=DEFAULT_BINARY)
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
parser.add_argument("--keep-work", action="store_true")
args = parser.parse_args()
if args.runs < 1:
parser.error("--runs must be positive")
if platform.system() != "Darwin":
parser.error("this benchmark currently targets the requested local macOS host")
binary = args.binary.resolve()
if not binary.is_file():
checked(["cargo", "build", "--release", "--locked"])
rustup = Path(shutil.which("rustup") or "")
if not rustup.is_file():
parser.error("rustup is required to use the current Rust 1.89 toolchain")
for command in ["git", "xcrun", "clang"]:
if not shutil.which(command):
parser.error(f"required host prerequisite is missing: {command}")
cpu_count = os.cpu_count() or 1
rustup_bin = rustup.parent
rustup_home = Path(os.environ.get("RUSTUP_HOME", Path.home() / ".rustup")).resolve()
args.output_dir.mkdir(parents=True, exist_ok=True)
temporary = Path(tempfile.mkdtemp(prefix="bot-forge-quality-bench-"))
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
log_dir = args.output_dir / f"quality-install-{stamp}-logs"
log_dir.mkdir(parents=True)
samples: list[Sample] = []
try:
prerequisite_versions = verify_rust_components(
sample_env(temporary / "prerequisite-check", rustup_bin, rustup_home, "serial", cpu_count)
)
for round_number in range(1, args.runs + 1):
methods = ["bot-forge", "serial"] if round_number % 2 else ["serial", "bot-forge"]
for order, method in enumerate(methods, 1):
print(f"[{round_number}/{args.runs}] {order}/2 {method}", flush=True)
samples.append(
run_sample(
method,
round_number,
order,
temporary,
binary,
rustup_bin,
rustup_home,
cpu_count,
log_dir,
)
)
forge_summary = summary(samples, "bot-forge")
serial_summary = summary(samples, "serial")
report = {
"schema_version": 1,
"generated_at": datetime.now(timezone.utc).isoformat(),
"host": {
"platform": platform.platform(),
"machine": platform.machine(),
"logical_cpus": cpu_count,
},
"toolchain": TOOLCHAIN,
"cargo_registry": RSPROXY_INDEX,
"targets": [name for name, _ in VERIFY],
"runs": args.runs,
"rust_components": {name: prerequisite_versions[name] for name, _ in VERIFY[:3]},
"summary": {"bot-forge": forge_summary, "serial": serial_summary},
"speedup": round(
serial_summary["median_seconds"] / forge_summary["median_seconds"], 3
),
"samples": [sample.__dict__ for sample in samples],
}
json_path = args.output_dir / f"quality-install-{stamp}.json"
md_path = args.output_dir / f"quality-install-{stamp}.md"
json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
md_path.write_text(markdown(report), encoding="utf-8")
print(md_path)
print(json_path)
return 0
except Exception as error:
print(f"benchmark failed: {error}", file=sys.stderr)
print(f"work directory: {temporary}", file=sys.stderr)
return 1
finally:
if not args.keep_work and len(samples) == args.runs * 2:
shutil.rmtree(temporary, ignore_errors=True)
if __name__ == "__main__":
raise SystemExit(main())