release: implement site registry v0.5

scripts/benchmark.py

@@ -0,0 +1,278 @@
+#!/usr/bin/env python3
+"""Run one registry operation and write a bounded, source-pinned benchmark report."""
+
+import argparse
+import hashlib
+import json
+import os
+import platform
+import resource
+import selectors
+import signal
+import sqlite3
+import subprocess
+import sys
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+
+MAX_MANIFEST_BYTES = 1024 * 1024
+MAX_OUTPUT_BYTES = 16 * 1024 * 1024
+
+
+def digest(data):
+    return hashlib.sha256(data).hexdigest()
+
+
+def file_digest(path):
+    value = hashlib.sha256()
+    with path.open("rb") as stream:
+        while chunk := stream.read(1024 * 1024):
+            value.update(chunk)
+    return value.hexdigest()
+
+
+def unique_object(pairs):
+    value = {}
+    for key, item in pairs:
+        if key in value:
+            raise ValueError("duplicate JSON key")
+        value[key] = item
+    return value
+
+
+def read_manifest(path):
+    if path.is_symlink() or not path.is_file():
+        raise ValueError("source manifest must be a regular file")
+    data = path.read_bytes()
+    if len(data) > MAX_MANIFEST_BYTES:
+        raise ValueError("source manifest exceeds 1 MiB")
+    manifest = json.loads(data, object_pairs_hook=unique_object)
+    required = ("schema", "source", "snapshot", "sha256", "bytes")
+    if not isinstance(manifest, dict) or any(key not in manifest for key in required):
+        raise ValueError("source manifest is incomplete")
+    if (not isinstance(manifest["sha256"], str)
+            or len(manifest["sha256"]) != 64
+            or any(char not in "0123456789abcdef" for char in manifest["sha256"])):
+        raise ValueError("source manifest has invalid SHA-256")
+    if not isinstance(manifest["bytes"], int) or manifest["bytes"] <= 0:
+        raise ValueError("source manifest has invalid byte count")
+    return {
+        "manifest_sha256": digest(data),
+        "schema": manifest["schema"],
+        "source": manifest["source"],
+        "snapshot": manifest["snapshot"],
+        "source_object_sha256": manifest["sha256"],
+        "compressed_bytes": manifest["bytes"],
+    }
+
+
+def database_counts(path):
+    if path is None or not path.exists():
+        return {"bytes": 0, "records": 0, "facts": 0}
+    if path.is_symlink() or not path.is_file():
+        raise ValueError("database must be a regular file")
+    connection = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
+    try:
+        page_count = connection.execute("PRAGMA page_count").fetchone()[0]
+        page_size = connection.execute("PRAGMA page_size").fetchone()[0]
+        tables = {
+            row[0]
+            for row in connection.execute(
+                "SELECT name FROM sqlite_schema WHERE type='table'"
+            )
+        }
+        records = connection.execute("SELECT count(*) FROM records").fetchone()[0] \
+            if "records" in tables else 0
+        facts = connection.execute("SELECT count(*) FROM facts").fetchone()[0] \
+            if "facts" in tables else 0
+        return {
+            "bytes": page_count * page_size,
+            "records": records,
+            "facts": facts,
+        }
+    finally:
+        connection.close()
+
+
+def bounded_file(path):
+    if path.stat().st_size > MAX_OUTPUT_BYTES:
+        return None
+    return path.read_bytes()
+
+
+def hardware():
+    cpu = "unknown"
+    cpuinfo = Path("/proc/cpuinfo")
+    if cpuinfo.is_file():
+        for line in cpuinfo.read_text(errors="replace").splitlines():
+            if line.startswith("model name"):
+                cpu = line.partition(":")[2].strip()
+                break
+    memory_kib = None
+    meminfo = Path("/proc/meminfo")
+    if meminfo.is_file():
+        first = meminfo.read_text(errors="replace").splitlines()[0].split()
+        if len(first) >= 2 and first[0] == "MemTotal:":
+            memory_kib = int(first[1])
+    return {
+        "operating_system": platform.system(),
+        "kernel": platform.release(),
+        "architecture": platform.machine(),
+        "cpu": cpu,
+        "logical_cpus": os.cpu_count(),
+        "memory_kib": memory_kib,
+    }
+
+
+def run_bounded(command, stdout_path, stderr_path):
+    process = subprocess.Popen(
+        command,
+        stdin=subprocess.DEVNULL,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        start_new_session=True,
+    )
+    streams = selectors.DefaultSelector()
+    streams.register(process.stdout, selectors.EVENT_READ, (stdout_path, 0))
+    streams.register(process.stderr, selectors.EVENT_READ, (stderr_path, 0))
+    exceeded = False
+    outputs = {
+        stdout_path: stdout_path.open("xb"),
+        stderr_path: stderr_path.open("xb"),
+    }
+    try:
+        while streams.get_map():
+            for key, _events in streams.select():
+                chunk = os.read(key.fileobj.fileno(), 64 * 1024)
+                path, written = key.data
+                if not chunk:
+                    streams.unregister(key.fileobj)
+                    key.fileobj.close()
+                    continue
+                accepted = min(len(chunk), MAX_OUTPUT_BYTES - written)
+                if accepted > 0:
+                    outputs[path].write(chunk[:accepted])
+                    written += accepted
+                    streams.modify(key.fileobj, selectors.EVENT_READ, (path, written))
+                if accepted != len(chunk) and not exceeded:
+                    exceeded = True
+                    try:
+                        os.killpg(process.pid, signal.SIGKILL)
+                    except ProcessLookupError:
+                        pass
+        return process.wait(), exceeded
+    finally:
+        streams.close()
+        for output in outputs.values():
+            output.close()
+
+
+def measured_artifacts(values):
+    result = []
+    labels = set()
+    for value in values:
+        label, separator, raw_path = value.partition("=")
+        if not separator or not label or label in labels:
+            raise ValueError("measured paths require unique LABEL=PATH values")
+        labels.add(label)
+        path = Path(raw_path)
+        if path.is_symlink() or not path.is_file():
+            raise ValueError("measured artifact must be a regular file")
+        result.append({
+            "label": label,
+            "bytes": path.stat().st_size,
+            "sha256": file_digest(path),
+        })
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--profile", choices=("small", "medium", "provider"), required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--source-manifest", type=Path, action="append", default=[])
+    parser.add_argument("--database", type=Path)
+    parser.add_argument("--expanded-bytes", type=int)
+    parser.add_argument("--measured-path", action="append", default=[])
+    parser.add_argument("command", nargs=argparse.REMAINDER)
+    args = parser.parse_args()
+    if not args.command or args.command[0] != "--" or len(args.command) == 1:
+        parser.error("terminate options with -- and provide a command")
+    command = args.command[1:]
+    if args.output.exists():
+        raise FileExistsError("benchmark output already exists")
+    if args.expanded_bytes is not None and args.expanded_bytes < 0:
+        raise ValueError("expanded bytes cannot be negative")
+    sources = [read_manifest(path) for path in args.source_manifest]
+    before = database_counts(args.database)
+    args.output.mkdir(parents=True)
+    stdout_path = args.output / "stdout"
+    stderr_path = args.output / "stderr"
+    started_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+    usage_before = resource.getrusage(resource.RUSAGE_CHILDREN)
+    started = time.monotonic()
+    exit_status, output_limit_exceeded = run_bounded(command, stdout_path, stderr_path)
+    wall_seconds = time.monotonic() - started
+    usage_after = resource.getrusage(resource.RUSAGE_CHILDREN)
+    after = database_counts(args.database)
+    stdout = bounded_file(stdout_path)
+    operation = None
+    if stdout is not None:
+        try:
+            operation = json.loads(stdout, object_pairs_hook=unique_object)
+        except (UnicodeDecodeError, json.JSONDecodeError, ValueError):
+            pass
+    fact_growth = after["facts"] - before["facts"]
+    report = {
+        "schema": "argand.site-benchmark/v1",
+        "profile": args.profile,
+        "started_at": started_at,
+        "hardware": hardware(),
+        "sources": sources,
+        "operation": {
+            "program": Path(command[0]).name,
+            "argument_vector_sha256": digest(b"\0".join(os.fsencode(item) for item in command)),
+            "exit_status": exit_status,
+            "output_limit_exceeded": output_limit_exceeded,
+            "stdout_sha256": file_digest(stdout_path),
+            "stderr_sha256": file_digest(stderr_path),
+            "result": operation,
+        },
+        "metrics": {
+            "wall_seconds": round(wall_seconds, 6),
+            "user_cpu_seconds": round(usage_after.ru_utime - usage_before.ru_utime, 6),
+            "system_cpu_seconds": round(usage_after.ru_stime - usage_before.ru_stime, 6),
+            "peak_rss_kib": usage_after.ru_maxrss,
+            "compressed_source_bytes": sum(source["compressed_bytes"] for source in sources),
+            "expanded_source_bytes": args.expanded_bytes,
+            "database_before": before,
+            "database_after": after,
+            "database_growth_bytes": after["bytes"] - before["bytes"],
+            "record_growth": after["records"] - before["records"],
+            "fact_growth": fact_growth,
+            "facts_per_second": round(fact_growth / wall_seconds, 3)
+            if fact_growth > 0 and wall_seconds > 0 else None,
+            "import_checkpoint_records": 256,
+        },
+        "artifacts": measured_artifacts(args.measured_path),
+    }
+    report_path = args.output / "REPORT.json"
+    with report_path.open("xb") as stream:
+        stream.write(json.dumps(report, indent=2, sort_keys=True).encode() + b"\n")
+        stream.flush()
+        os.fsync(stream.fileno())
+    directory = os.open(args.output, os.O_RDONLY | os.O_DIRECTORY)
+    try:
+        os.fsync(directory)
+    finally:
+        os.close(directory)
+    return exit_status
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except (FileExistsError, OSError, ValueError) as error:
+        print(error, file=sys.stderr)
+        sys.exit(2)