Nexus/tests/load_test.py

"""
Nexus 负载压测脚本
测试: /health, /health/services, /api/servers/stats, /api/agent/heartbeat, /api/servers/exec

用法:
  NEXUS_API_KEY=your-key python tests/load_test.py [--concurrency 50] [--duration 30]
  NEXUS_TEST_BASE=http://127.0.0.1:8600  (可选)

注意 (BL-06): /api/agent/heartbeat 仅接受 per-server agent_api_key，不再接受 global API_KEY。
  心跳压测需设置该服务器的 per-server key（与 SERVER_ID 对应），或跳过 heartbeat worker。
"""
import asyncio
import httpx
import time
import sys
import os
import argparse

BASE = os.environ.get("NEXUS_TEST_BASE", "http://127.0.0.1:8600")
API_KEY = os.environ.get("NEXUS_API_KEY", "")

# 取一台在线服务器做 exec 测试
SERVER_ID = 1  # 默认，会动态查找

stats = {"total": 0, "ok": 0, "fail": 0, "times": [], "errors": []}


async def hit(client: httpx.AsyncClient, name: str, method: str, path: str, body=None, need_key=True):
    headers = {}
    if need_key and API_KEY:
        headers["X-API-Key"] = API_KEY
    t0 = time.perf_counter()
    try:
        if method == "GET":
            r = await client.get(f"{BASE}{path}", headers=headers, timeout=30)
        elif method == "POST":
            r = await client.post(f"{BASE}{path}", json=body, headers=headers, timeout=30)
        elapsed = time.perf_counter() - t0
        if r.status_code == 200:
            stats["ok"] += 1
        else:
            stats["fail"] += 1
            stats["errors"].append(f"[{r.status_code}] {name} {path}")
    except Exception as e:
        elapsed = time.perf_counter() - t0
        stats["fail"] += 1
        stats["errors"].append(f"[EXCEPTION] {name}: {str(e)[:80]}")
    stats["total"] += 1
    stats["times"].append(elapsed)


async def health_worker(client: httpx.AsyncClient, duration: float):
    """轻量级 health check 持续请求"""
    end = time.time() + duration
    while time.time() < end:
        await hit(client, "health", "GET", "/health", need_key=False)
        await asyncio.sleep(0.01)


async def health_services_worker(client: httpx.AsyncClient, duration: float):
    """health/services 持续请求"""
    end = time.time() + duration
    while time.time() < end:
        await hit(client, "health/services", "GET", "/health/services", need_key=False)
        await asyncio.sleep(0.05)


async def stats_worker(client: httpx.AsyncClient, duration: float):
    """仪表盘 stats (Redis 缓存)"""
    end = time.time() + duration
    while time.time() < end:
        await hit(client, "stats", "GET", "/api/servers/stats")
        await asyncio.sleep(0.1)


async def heartbeat_worker(client: httpx.AsyncClient, duration: float):
    """Agent 心跳压测"""
    end = time.time() + duration
    while time.time() < end:
        await hit(client, "heartbeat", "POST", "/api/agent/heartbeat", body={
            "cpu": 25.0, "memory": 60.5, "disk": 45.2,
            "load_1": 0.5, "load_5": 0.3, "load_15": 0.2,
            "uptime": 86400, "os": "Ubuntu 24.04"
        }, need_key=False)
        await asyncio.sleep(0.2)


async def exec_worker(client: httpx.AsyncClient, duration: float, server_id: int):
    """群发命令压测 (走 ThreadPoolExecutor SSH)"""
    end = time.time() + duration
    while time.time() < end:
        await hit(client, "exec", "POST", f"/api/servers/{server_id}/exec", body={
            "command": "echo 'load-test' && date"
        })
        await asyncio.sleep(0.5)


def print_report(concurrency: int, duration: float):
    times = sorted(stats["times"]) if stats["times"] else [0]
    n = len(times)
    p50 = times[int(n * 0.50)] if n > 0 else 0
    p95 = times[min(int(n * 0.95), n - 1)] if n > 0 else 0
    p99 = times[min(int(n * 0.99), n - 1)] if n > 0 else 0
    rps = stats["total"] / duration if duration > 0 else 0
    err_rate = stats["fail"] / stats["total"] * 100 if stats["total"] > 0 else 0

    print(f"""
{'='*60}
  Nexus 压测报告
{'='*60}
  并发:        {concurrency}
  持续时间:    {duration}s
  总请求:      {stats['total']}
  成功:        {stats['ok']}
  失败:        {stats['fail']}
  错误率:      {err_rate:.1f}%
  RPS:         {rps:.1f}
  延迟 P50:    {p50*1000:.1f}ms
  延迟 P95:    {p95*1000:.1f}ms
  延迟 P99:    {p99*1000:.1f}ms
  最慢:        {max(times)*1000:.1f}ms
  最快:        {min(times)*1000:.1f}ms
{'='*60}
""")
    if stats["errors"]:
        print("错误采样 (前10):")
        for e in stats["errors"][:10]:
            print(f"  {e}")


async def main():
    global BASE, API_KEY, SERVER_ID, stats
    parser = argparse.ArgumentParser()
    parser.add_argument("--concurrency", type=int, default=50)
    parser.add_argument("--duration", type=int, default=30)
    parser.add_argument("--no-exec", action="store_true", help="跳过 SSH exec 测试")
    args = parser.parse_args()

    global BASE, API_KEY
    BASE = os.environ.get("NEXUS_TEST_BASE", "http://127.0.0.1:8600")
    API_KEY = os.environ.get("NEXUS_API_KEY", "")
    if not API_KEY:
        print("WARN: NEXUS_API_KEY unset — stats/exec workers may be skipped", file=sys.stderr)

    concurrency = args.concurrency
    duration = args.duration
    print(f"压测配置: BASE={BASE}, 并发={concurrency}, 时长={duration}s")

    # 动态获取一台在线服务器
    async with httpx.AsyncClient(timeout=10) as client:
        if not args.no_exec and API_KEY:
            try:
                r = await client.get(f"{BASE}/api/servers/", headers={"X-API-Key": API_KEY})
                servers = r.json() if r.status_code == 200 else []
                if isinstance(servers, dict):
                    servers = servers.get("data", [])
                online = [s for s in servers if s.get("is_online")]
                if online:
                    SERVER_ID = online[0]["id"]
                    print(f"exec 目标服务器: id={SERVER_ID} name={online[0].get('name', '?')}")
                else:
                    print("无在线服务器，跳过 exec 测试")
                    args.no_exec = True
            except Exception as e:
                print(f"获取服务器列表失败: {e}，跳过 exec 测试")
                args.no_exec = True
        else:
            args.no_exec = True

    stats = {"total": 0, "ok": 0, "fail": 0, "times": [], "errors": []}

    async with httpx.AsyncClient(limits=httpx.Limits(max_connections=concurrency + 100, max_keepalive_connections=concurrency)) as client:
        tasks = []

        # 分配权重: health 40%, stats 20%, health-services 10%, heartbeat 20%, exec 10%
        n_health = max(1, int(concurrency * 0.4))
        n_stats = max(1, int(concurrency * 0.2))
        n_hs = max(1, int(concurrency * 0.1))
        n_hb = max(1, int(concurrency * 0.2))
        n_exec = max(1, int(concurrency * 0.1))

        print(f"Worker: health={n_health} stats={n_stats} health-svc={n_hs} heartbeat={n_hb} exec={n_exec}")

        for _ in range(n_health):
            tasks.append(health_worker(client, duration))
        for _ in range(n_stats):
            tasks.append(stats_worker(client, duration))
        for _ in range(n_hs):
            tasks.append(health_services_worker(client, duration))
        for _ in range(n_hb):
            tasks.append(heartbeat_worker(client, duration))
        if not args.no_exec:
            for _ in range(n_exec):
                tasks.append(exec_worker(client, duration, SERVER_ID))

        print(f"启动 {len(tasks)} 个 worker...")
        t0 = time.time()
        await asyncio.gather(*tasks, return_exceptions=True)
        elapsed = time.time() - t0
        print_report(concurrency, elapsed)


if __name__ == "__main__":
    asyncio.run(main())