Nexus Agent
eab35a35be
移除 Agent 认证与安装链路的 global API_KEY 回退，收口 risk-acceptance 接受项 2；巡检补修心跳重复 get_server 并加 assert_called_once 回归测试。 Co-authored-by: Cursor <cursoragent@cursor.com>
205 lines
7.5 KiB
Python
205 lines
7.5 KiB
Python
"""
|
|
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())
|