multi-agent-mux/.agents/skills/lib_py/delegation/bench.py

"""Benchmark harness for the persistent-connection prototype.

Two transports, deliberately separated:

  --transport fake    a stub client with a synthetic connect latency. Fully
                      deterministic, touches no network, and is what CI asserts
                      on: N publishes cost N connections one-shot and exactly 1
                      persistent.
  --transport broker  the real broker named by the job record. Informational
                      only: the PoC default is a public broker whose CONNACK
                      latency is not ours to control, so no gate depends on it.

Usage:
  python -m lib_py.delegation.bench --n 50 --transport fake [--connect-latency-ms 50]
                                    [--job <id>] [--registry-dir DIR] [--out FILE]
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from typing import List, Optional

from lib_py.delegation import mqtt_common, publish_event
from lib_py.delegation.publisher import OneShotPublisher, PersistentPublisher


class FakeInfo(object):
    """Stand-in for paho's MQTTMessageInfo, including its failure contract.

    paho carries an ``rc`` and raises RuntimeError from ``wait_for_publish``
    when the publish could not be handed to a live connection. The earlier
    version of this fake had neither, which is precisely why a real defect in
    that path stayed invisible while every test passed.
    """

    def __init__(self, published=True, rc=0):
        self._published = published
        self.rc = rc

    def wait_for_publish(self, timeout=None):
        if self.rc:
            raise RuntimeError(
                "Message publish failed: The client is not currently connected.")
        return None

    def is_published(self):
        return self._published


class FakeClient(object):
    """Minimal paho stand-in: counts connects and charges a synthetic latency."""

    def __init__(self, connect_latency=0.0, counters=None):
        self.connect_latency = connect_latency
        self.counters = counters if counters is not None else {}
        self.on_connect = None

    def connect(self, host, port, keepalive=60):
        if self.connect_latency:
            time.sleep(self.connect_latency)
        self.counters["connect"] = self.counters.get("connect", 0) + 1
        if self.on_connect:
            self.on_connect(self, None, None, 0, None)

    def loop_start(self):
        return None

    def loop_stop(self):
        return None

    def reconnect_delay_set(self, min_delay=1, max_delay=16):
        return None

    def publish(self, topic, payload=None, qos=0, retain=False):
        self.counters["publish"] = self.counters.get("publish", 0) + 1
        return FakeInfo()

    def disconnect(self):
        return None


def fake_factory(connect_latency, counters):
    def factory(_role, _config):
        return FakeClient(connect_latency=connect_latency, counters=counters)
    return factory


def _percentile(values: List[float], pct: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    idx = int(round((pct / 100.0) * (len(ordered) - 1)))
    return ordered[idx]


def _run(publisher, prepared, n: int) -> dict:
    latencies = []
    failures = 0
    for _ in range(n):
        start = time.monotonic()
        ok, _err = publisher.publish_with_retry(prepared, attempts=1)
        latencies.append((time.monotonic() - start) * 1000.0)
        if not ok:
            failures += 1
    return {
        "publishes": n,
        "connects": getattr(publisher, "connect_count", None),
        "failures": failures,
        "latency_ms": {
            "p50": round(_percentile(latencies, 50), 3),
            "p95": round(_percentile(latencies, 95), 3),
            "max": round(max(latencies), 3) if latencies else 0.0,
        },
    }


def run_benchmark(n=50, transport="fake", connect_latency_ms=50.0,
                  job_id=None, registry_dir=None) -> dict:
    registry_dir = registry_dir or mqtt_common.DEFAULT_REGISTRY_DIR

    if transport == "fake":
        prepared = publish_event.PreparedEvent(
            job_id="benchjob", event="progress", detail="bench", seq=0,
            topic="python/mqtt/jobs/benchjob/events", retain=False,
            payload={"seq": 0}, body=b'{"seq": 0}',
            config=mqtt_common.broker_config_from_env(),
        )
        latency = connect_latency_ms / 1000.0
        one_counters, persist_counters = {}, {}
        one = OneShotPublisher(client_factory=fake_factory(latency, one_counters))
        persist = PersistentPublisher(client_factory=fake_factory(latency, persist_counters))
        try:
            result = {
                "transport": "fake",
                "connect_latency_ms": connect_latency_ms,
                "one_shot": _run(one, prepared, n),
                "persistent": _run(persist, prepared, n),
            }
        finally:
            persist.close()
        result["one_shot"]["client_connects"] = one_counters.get("connect", 0)
        result["persistent"]["client_connects"] = persist_counters.get("connect", 0)
        return result

    if not job_id:
        raise ValueError("--transport broker requires --job")
    prepared = publish_event.prepare_event(
        job_id=job_id, event="progress", detail="bench", registry_dir=registry_dir,
    )
    one = OneShotPublisher()
    persist = PersistentPublisher()
    try:
        return {
            "transport": "broker",
            "broker": "{}:{}".format(prepared.config.host, prepared.config.port),
            "one_shot": _run(one, prepared, n),
            "persistent": _run(persist, prepared, n),
        }
    finally:
        persist.close()


def main(argv: Optional[List[str]] = None) -> int:
    mqtt_common.ensure_env_loaded()
    parser = argparse.ArgumentParser(prog="lib_py.delegation.bench")
    parser.add_argument("--n", type=int, default=50)
    parser.add_argument("--transport", default="fake", choices=["fake", "broker"])
    parser.add_argument("--connect-latency-ms", type=float, default=50.0)
    parser.add_argument("--job", default=None)
    parser.add_argument("--registry-dir", default=None)
    parser.add_argument("--out", default=None)
    args = parser.parse_args(argv)

    result = run_benchmark(
        n=args.n, transport=args.transport,
        connect_latency_ms=args.connect_latency_ms,
        job_id=args.job, registry_dir=args.registry_dir,
    )
    text = json.dumps(result, indent=2, sort_keys=True)
    if args.out:
        with open(args.out, "w", encoding="utf-8") as fh:
            fh.write(text + "\n")
    print(text)
    return 0


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