@agentunion/kite 1.0.7 → 1.3.0

package/core/event_hub/bench.py

@@ -1,459 +0,0 @@
-"""
-Event Hub stress test (standalone).
-
-Usage:  python -m core.event_hub.bench
-        (from Kite root directory)
-
-Phase 1: 3-channel mixed stress test (10 minutes)
-  CH1 — high-freq small messages (500 evt/s)
-  CH2 — medium-freq large messages (50 evt/s, 10KB)
-  CH3 — bursty traffic (200-event burst every 2s)
-
-Phase 2: 3 extreme tests
-  EX1 — max burst (ramp until failure)
-  EX2 — max concurrent connections
-  EX3 — max message size
-"""
-
-import asyncio
-import json
-import os
-import socket
-import statistics
-import sys
-import threading
-import time
-import uuid
-from datetime import datetime, timezone
-
-_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-if _root not in sys.path:
-    sys.path.insert(0, _root)
-
-import psutil
-import uvicorn
-import websockets
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-
-from core.event_hub.hub import EventHub
-
-
-# ── Minimal Event Hub server (no auth) ──
-
-def _create_app(hub: EventHub) -> FastAPI:
-    app = FastAPI()
-
-    @app.websocket("/ws")
-    async def ws(ws: WebSocket):
-        mid = ws.query_params.get("id", f"anon_{id(ws)}")
-        await ws.accept()
-        hub.add_connection(mid, ws)
-        try:
-            while True:
-                raw = await ws.receive_text()
-                await hub.handle_message(mid, ws, raw)
-        except (WebSocketDisconnect, Exception):
-            pass
-        finally:
-            hub.remove_connection(mid)
-
-    return app
-
-
-def _free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return s.getsockname()[1]
-
-
-# ── Client ──
-
-class C:
-    """Lightweight WS client."""
-    def __init__(self, url: str, name: str):
-        self.url, self.name, self.ws = url, name, None
-
-    async def connect(self):
-        self.ws = await websockets.connect(f"{self.url}?id={self.name}", max_size=None)
-
-    async def sub(self, patterns):
-        await self.ws.send(json.dumps({"type": "subscribe", "events": patterns}))
-
-    async def pub(self, event: str, data: dict) -> str:
-        eid = str(uuid.uuid4())
-        await self.ws.send(json.dumps({
-            "type": "event", "event_id": eid, "event": event,
-            "source": self.name, "timestamp": datetime.now(timezone.utc).isoformat(),
-            "data": data,
-        }))
-        return eid
-
-    async def recv(self, timeout=5.0):
-        try:
-            return json.loads(await asyncio.wait_for(self.ws.recv(), timeout=timeout))
-        except Exception:
-            return None
-
-    async def drain(self, timeout=0.05):
-        while await self.recv(timeout=timeout) is not None:
-            pass
-
-    async def close(self):
-        if self.ws:
-            await self.ws.close()
-
-
-# ── Stats collector ──
-
-class Stats:
-    def __init__(self):
-        self.sent = 0
-        self.received = 0
-        self.acked = 0
-        self.latencies: list[float] = []
-        self.lock = asyncio.Lock()
-
-    async def record_latency(self, ts_iso: str):
-        try:
-            sent = datetime.fromisoformat(ts_iso)
-            ms = (datetime.now(timezone.utc) - sent).total_seconds() * 1000
-            async with self.lock:
-                self.latencies.append(ms)
-        except Exception:
-            pass
-
-    def summary(self) -> dict:
-        lat = sorted(self.latencies) if self.latencies else [0]
-        return {
-            "sent": self.sent,
-            "recv": self.received,
-            "acked": self.acked,
-            "loss": self.sent - self.received,
-            "lat_avg": round(statistics.mean(lat), 1),
-            "lat_p50": round(statistics.median(lat), 1),
-            "lat_p95": round(lat[int(len(lat) * 0.95)], 1),
-            "lat_p99": round(lat[int(len(lat) * 0.99)], 1),
-        }
-
-
-# ── Benchmark ──
-
-class Benchmark:
-    def __init__(self):
-        self.hub = EventHub()
-        self.port = _free_port()
-        self.ws_url = f"ws://127.0.0.1:{self.port}/ws"
-        self._server: uvicorn.Server | None = None
-        self.proc = psutil.Process()
-
-    def _start_server(self):
-        app = _create_app(self.hub)
-        cfg = uvicorn.Config(app, host="127.0.0.1", port=self.port, log_level="warning")
-        self._server = uvicorn.Server(cfg)
-        threading.Thread(target=self._server.run, daemon=True).start()
-        deadline = time.time() + 5
-        while time.time() < deadline:
-            if self._server.started:
-                return
-            time.sleep(0.05)
-        raise RuntimeError("Server failed to start")
-
-    def _res_snapshot(self) -> dict:
-        return {
-            "cpu": self.proc.cpu_percent(),
-            "mem_mb": round(self.proc.memory_info().rss / 1024 / 1024, 1),
-        }
-
-    async def _make_channel(self, pub_name: str, sub_name: str, pattern: str):
-        """Create a pub/sub pair, subscribe, return (pub, sub, stats)."""
-        pub = C(self.ws_url, pub_name)
-        sub = C(self.ws_url, sub_name)
-        await pub.connect()
-        await sub.connect()
-        await sub.sub([pattern])
-        await asyncio.sleep(0.1)
-        return pub, sub, Stats()
-
-    # ── ACK drainer (keeps publisher WS buffer clear) ──
-
-    async def _drainer(self, pub: C, stats: Stats, stop: asyncio.Event):
-        while not stop.is_set():
-            msg = await pub.recv(timeout=0.3)
-            if msg and msg.get("type") == "ack":
-                stats.acked += 1
-
-    # ── Subscriber receiver (shared by all channels) ──
-
-    async def _receiver(self, sub: C, stats: Stats, stop: asyncio.Event):
-        while not stop.is_set():
-            msg = await sub.recv(timeout=0.3)
-            if not msg:
-                continue
-            if msg.get("type") == "event":
-                stats.received += 1
-                ts = msg.get("timestamp", "")
-                if ts:
-                    await stats.record_latency(ts)
-
-    # ── CH1: high-freq small messages ──
-
-    async def _ch1_sender(self, pub: C, stats: Stats, stop: asyncio.Event):
-        """500 evt/s, tiny payload."""
-        while not stop.is_set():
-            await pub.pub("ch1.fast", {"seq": stats.sent})
-            stats.sent += 1
-            await asyncio.sleep(0.002)
-
-    # ── CH2: medium-freq large messages ──
-
-    async def _ch2_sender(self, pub: C, stats: Stats, stop: asyncio.Event):
-        """50 evt/s, 10KB payload."""
-        payload = "X" * 10240
-        while not stop.is_set():
-            await pub.pub("ch2.large", {"p": payload})
-            stats.sent += 1
-            await asyncio.sleep(0.02)
-
-    # ── CH3: bursty traffic ──
-
-    async def _ch3_sender(self, pub: C, stats: Stats, stop: asyncio.Event):
-        """200-event burst every 2s."""
-        while not stop.is_set():
-            for _ in range(200):
-                await pub.pub("ch3.burst", {"b": stats.sent})
-                stats.sent += 1
-            await asyncio.sleep(2)
-
-    # ── Phase 1: 10-min mixed stress ──
-
-    async def phase1(self, duration=600):
-        print("=" * 60)
-        print(f"PHASE 1: 3-channel mixed stress test ({duration}s)")
-        print("=" * 60)
-
-        p1, s1, st1 = await self._make_channel("ch1_pub", "ch1_sub", "ch1.*")
-        p2, s2, st2 = await self._make_channel("ch2_pub", "ch2_sub", "ch2.*")
-        p3, s3, st3 = await self._make_channel("ch3_pub", "ch3_sub", "ch3.*")
-        stop = asyncio.Event()
-
-        tasks = [
-            asyncio.create_task(self._ch1_sender(p1, st1, stop)),
-            asyncio.create_task(self._drainer(p1, st1, stop)),
-            asyncio.create_task(self._receiver(s1, st1, stop)),
-            asyncio.create_task(self._ch2_sender(p2, st2, stop)),
-            asyncio.create_task(self._drainer(p2, st2, stop)),
-            asyncio.create_task(self._receiver(s2, st2, stop)),
-            asyncio.create_task(self._ch3_sender(p3, st3, stop)),
-            asyncio.create_task(self._drainer(p3, st3, stop)),
-            asyncio.create_task(self._receiver(s3, st3, stop)),
-        ]
-
-        self.proc.cpu_percent()  # prime
-        start = time.time()
-        interval = 30  # report every 30s
-        next_report = start + interval
-
-        while time.time() - start < duration:
-            await asyncio.sleep(1)
-            if time.time() >= next_report:
-                elapsed = round(time.time() - start)
-                res = self._res_snapshot()
-                hub_stats = self.hub._counters_dict()
-                print(f"\n  [{elapsed:>4}s] cpu={res['cpu']:.0f}%  mem={res['mem_mb']}MB  "
-                      f"hub_recv={hub_stats['events_received']}  hub_route={hub_stats['events_routed']}")
-                for name, st in [("CH1-fast", st1), ("CH2-large", st2), ("CH3-burst", st3)]:
-                    s = st.summary()
-                    print(f"    {name}: sent={s['sent']}  recv={s['recv']}  loss={s['loss']}  "
-                          f"avg={s['lat_avg']}ms  p95={s['lat_p95']}ms  p99={s['lat_p99']}ms")
-                next_report = time.time() + interval
-
-        stop.set()
-        for t in tasks:
-            t.cancel()
-        await asyncio.gather(*tasks, return_exceptions=True)
-
-        # Final report
-        print(f"\n  FINAL ({duration}s):")
-        res = self._res_snapshot()
-        print(f"  cpu={res['cpu']:.0f}%  mem={res['mem_mb']}MB")
-        for name, st in [("CH1-fast", st1), ("CH2-large", st2), ("CH3-burst", st3)]:
-            s = st.summary()
-            print(f"    {name}: sent={s['sent']}  recv={s['recv']}  loss={s['loss']}  "
-                  f"avg={s['lat_avg']}ms  p50={s['lat_p50']}ms  p95={s['lat_p95']}ms  p99={s['lat_p99']}ms")
-
-        for c in [p1, s1, p2, s2, p3, s3]:
-            await c.close()
-        print()
-
-    # ── Phase 2: Extreme tests ──
-
-    async def extreme1_max_burst(self):
-        """Ramp burst size until loss or timeout."""
-        print("=" * 60)
-        print("EXTREME 1: Max burst (ramp until loss)")
-        print("=" * 60)
-
-        pub, sub, _ = await self._make_channel("ex1_pub", "ex1_sub", "ex1.*")
-
-        for size in [1000, 5000, 10000, 20000, 50000]:
-            recvd = 0
-            stop = asyncio.Event()
-
-            async def _recv():
-                nonlocal recvd
-                while not stop.is_set():
-                    msg = await sub.recv(timeout=0.3)
-                    if msg and msg.get("type") == "event":
-                        recvd += 1
-
-            task = asyncio.create_task(_recv())
-            start = time.time()
-            for i in range(size):
-                await pub.pub("ex1.burst", {"i": i})
-            send_time = time.time() - start
-
-            # Wait for delivery
-            await asyncio.sleep(max(3, size / 2000))
-            stop.set()
-            await task
-            await pub.drain()
-            await sub.drain()
-
-            rate = size / send_time if send_time > 0 else 0
-            loss = size - recvd
-            loss_pct = loss / size * 100
-            res = self._res_snapshot()
-            print(f"  burst={size:>6}: recv={recvd}  loss={loss}({loss_pct:.1f}%)  "
-                  f"rate={rate:.0f} evt/s  cpu={res['cpu']:.0f}%  mem={res['mem_mb']}MB")
-
-            if loss_pct > 5:
-                print(f"  >> Loss exceeded 5%, stopping ramp")
-                break
-
-        await pub.close()
-        await sub.close()
-        print()
-
-    async def extreme2_max_connections(self):
-        """Ramp concurrent connections, each publishing 1 evt/s."""
-        print("=" * 60)
-        print("EXTREME 2: Max concurrent connections")
-        print("=" * 60)
-
-        # One global subscriber
-        gsub = C(self.ws_url, "ex2_gsub")
-        await gsub.connect()
-        await gsub.sub(["ex2.*"])
-        await asyncio.sleep(0.1)
-
-        for n_conn in [10, 50, 100, 200, 500]:
-            clients = []
-            try:
-                for i in range(n_conn):
-                    c = C(self.ws_url, f"ex2_c{i}")
-                    await c.connect()
-                    clients.append(c)
-            except Exception as e:
-                print(f"  {len(clients)} connections: FAILED to open more ({e})")
-                break
-
-            # Each client sends 1 event
-            recvd = 0
-            stop = asyncio.Event()
-
-            async def _recv():
-                nonlocal recvd
-                while not stop.is_set():
-                    msg = await gsub.recv(timeout=0.3)
-                    if msg and msg.get("type") == "event":
-                        recvd += 1
-
-            task = asyncio.create_task(_recv())
-            start = time.time()
-            for c in clients:
-                await c.pub("ex2.ping", {"from": c.name})
-            send_time = time.time() - start
-
-            await asyncio.sleep(max(2, n_conn / 200))
-            stop.set()
-            await task
-
-            res = self._res_snapshot()
-            print(f"  {n_conn:>4} conns: recv={recvd}/{n_conn}  "
-                  f"send_time={send_time:.2f}s  cpu={res['cpu']:.0f}%  mem={res['mem_mb']}MB")
-
-            for c in clients:
-                await c.close()
-
-        await gsub.close()
-        print()
-
-    async def extreme3_max_message_size(self):
-        """Ramp message size until failure."""
-        print("=" * 60)
-        print("EXTREME 3: Max message size")
-        print("=" * 60)
-
-        pub, sub, _ = await self._make_channel("ex3_pub", "ex3_sub", "ex3.*")
-
-        for size_kb in [100, 500, 1000, 2000, 5000]:
-            payload = "Z" * (size_kb * 1024)
-            latencies = []
-            ok = True
-
-            for _ in range(5):
-                try:
-                    await pub.pub("ex3.big", {"p": payload})
-                except Exception as e:
-                    print(f"  {size_kb:>5} KB: SEND FAILED ({e})")
-                    ok = False
-                    break
-
-            if not ok:
-                break
-
-            await asyncio.sleep(max(1, size_kb / 500))
-            # Collect
-            for _ in range(20):
-                msg = await sub.recv(timeout=0.3)
-                if msg and msg.get("type") == "event":
-                    ts = msg.get("timestamp", "")
-                    try:
-                        sent = datetime.fromisoformat(ts)
-                        latencies.append((datetime.now(timezone.utc) - sent).total_seconds() * 1000)
-                    except Exception:
-                        pass
-                if msg is None:
-                    break
-
-            await pub.drain()
-            res = self._res_snapshot()
-            if latencies:
-                avg = statistics.mean(latencies)
-                print(f"  {size_kb:>5} KB: recv={len(latencies)}/5  "
-                      f"avg={avg:.1f}ms  cpu={res['cpu']:.0f}%  mem={res['mem_mb']}MB")
-            else:
-                print(f"  {size_kb:>5} KB: recv=0/5  cpu={res['cpu']:.0f}%  mem={res['mem_mb']}MB")
-
-        await pub.close()
-        await sub.close()
-        print()
-
-    # ── Entry ──
-
-    async def run_all(self):
-        self._start_server()
-        print(f"Event Hub started on port {self.port}\n")
-        try:
-            await self.phase1(duration=600)
-            await self.extreme1_max_burst()
-            await self.extreme2_max_connections()
-            await self.extreme3_max_message_size()
-            print("All tests complete.")
-        finally:
-            if self._server:
-                self._server.should_exit = True
-
-
-if __name__ == "__main__":
-    asyncio.run(Benchmark().run_all())