120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
"""Private benchmark wrapper for counting browser-harness calls.
|
|
|
|
The wrapper is deliberately a transport-only probe. It records method names and
|
|
counts, never CDP arguments, URLs, DOM values, page text, or screenshots. The
|
|
selected worker is executed unchanged after the helper module is wrapped.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import runpy
|
|
import sys
|
|
import tempfile
|
|
from collections import Counter
|
|
from typing import Any
|
|
|
|
|
|
def _take_flag(argv: list[str], flag: str) -> tuple[str | None, list[str]]:
|
|
remaining: list[str] = []
|
|
value: str | None = None
|
|
index = 0
|
|
while index < len(argv):
|
|
item = argv[index]
|
|
if item == flag:
|
|
if value is not None or index + 1 >= len(argv):
|
|
raise ValueError(f"{flag} must be supplied once")
|
|
value = argv[index + 1]
|
|
index += 2
|
|
continue
|
|
remaining.append(item)
|
|
index += 1
|
|
return value, remaining
|
|
|
|
|
|
def _method_name(value: Any) -> str:
|
|
if isinstance(value, str) and value and len(value) <= 120 and all(
|
|
character.isalnum() or character in "._-" for character in value
|
|
):
|
|
return value
|
|
return "unknown"
|
|
|
|
|
|
def _wrap_helpers(counts: dict[str, Counter[str]] | None = None) -> dict[str, Counter[str]]:
|
|
helpers = importlib.import_module("browser_harness.helpers")
|
|
counts = counts or {"cdp": Counter(), "ipc": Counter()}
|
|
|
|
original_cdp = getattr(helpers, "cdp", None)
|
|
if callable(original_cdp):
|
|
def counted_cdp(*args: Any, **kwargs: Any) -> Any:
|
|
method = _method_name(args[0] if args else None)
|
|
counts["cdp"][method] += 1
|
|
return original_cdp(*args, **kwargs)
|
|
helpers.cdp = counted_cdp
|
|
|
|
original_send = getattr(helpers, "_send", None)
|
|
if callable(original_send):
|
|
def counted_send(*args: Any, **kwargs: Any) -> Any:
|
|
# Do not inspect the request. The IPC method is intentionally the
|
|
# fixed helper method name rather than an untrusted operation value.
|
|
counts["ipc"]["_send"] += 1
|
|
return original_send(*args, **kwargs)
|
|
helpers._send = counted_send
|
|
return counts
|
|
|
|
|
|
def _write_counts(path: str, counts: dict[str, Counter[str]]) -> None:
|
|
payload = {
|
|
"cdp": dict(sorted(counts["cdp"].items())),
|
|
"ipc": dict(sorted(counts["ipc"].items())),
|
|
}
|
|
destination = Path(path)
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
fd, temporary = tempfile.mkstemp(prefix="performance-count-", dir=str(destination.parent))
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
|
json.dump(payload, stream, separators=(",", ":"), sort_keys=True)
|
|
stream.write("\n")
|
|
os.chmod(temporary, 0o600)
|
|
os.replace(temporary, destination)
|
|
os.chmod(destination, 0o600)
|
|
finally:
|
|
try:
|
|
os.unlink(temporary)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
raw = list(sys.argv[1:] if argv is None else argv)
|
|
worker_path, raw = _take_flag(raw, "--performance-worker")
|
|
count_path, worker_argv = _take_flag(raw, "--performance-count-file")
|
|
if not worker_path or not count_path:
|
|
raise SystemExit("performance_probe requires --performance-worker and --performance-count-file")
|
|
if not Path(worker_path).is_file():
|
|
raise SystemExit("performance worker path is not a file")
|
|
counts: dict[str, Counter[str]] = {"cdp": Counter(), "ipc": Counter()}
|
|
try:
|
|
sys.argv = [worker_path, *worker_argv]
|
|
namespace = runpy.run_path(worker_path, run_name="performance_worker")
|
|
backend = namespace.get("RealBrowserBackend")
|
|
main_function = namespace.get("main")
|
|
if backend is None or not callable(main_function):
|
|
raise SystemExit("performance worker does not expose its expected entry points")
|
|
original_load = backend._load_harness
|
|
def load_and_probe(instance: Any) -> Any:
|
|
result = original_load(instance)
|
|
_wrap_helpers(counts)
|
|
return result
|
|
backend._load_harness = load_and_probe
|
|
main_function(worker_argv)
|
|
finally:
|
|
_write_counts(count_path, counts)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|