robotex/tests/loadtest.py

"""Нагрузочный тест: N параллельных запросов на каждом уровне (каждый — в своей сессии) + пиковое
потребление контейнера воркера по `docker stats`.

	docker compose --profile loadtest up -d
	python tests/loadtest.py --api-key rbx_... --levels 1 5 10 15            # /v1/solve
	python tests/loadtest.py --api-key rbx_... --mode v2 --levels 4 8        # /v2/tasks + опрос

Помимо status ответа, неуспех раскладывается по причинам: нет слотов, застряли на Cloudflare,
упало действие, исключение до действий.
"""
import argparse
import asyncio
import json
import os
import statistics
import subprocess
import threading
import time

import httpx

CF_MARKERS = ("challenges.cloudflare.com", "Just a moment", "cf-turnstile")


def build_payload(args) -> dict:
	return {
		"url": args.url,
		"proxy": args.proxy,
		"timeout": 120,
		"html": True,
		"captcha_type": args.captcha_type,
		"actions": [{"type": "wait_for", "selector": args.selector, "timeout": 60}],
	}


def classify(resp: httpx.Response | Exception) -> str:
	if isinstance(resp, Exception):
		return f"exc:{type(resp).__name__}"
	data = resp.json()
	code = data.get("code")
	if code == "no_slots":
		return "no_slots"
	if code == "challenge_failed":
		return "cf_blocked"
	if code == "action_failed":
		html = data.get("html") or ""
		return "cf_blocked" if any(m in html for m in CF_MARKERS) else "action_failed"
	if data.get("status") != "ok":
		return f"error:{str(data.get('error'))[:60]}"
	return "ok"


class StatsSampler(threading.Thread):
	"""Раз в ~2с снимает docker stats контейнера, копит пики CPU/памяти и число процессов."""

	def __init__(self, container: str):
		super().__init__(daemon=True)
		self.container = container
		self.stop_event = threading.Event()
		self.cpu = 0.0
		self.mem_mb = 0.0
		self.pids = 0

	@staticmethod
	def _to_mb(value: str) -> float:
		units = {"KiB": 1 / 1024, "MiB": 1, "GiB": 1024, "kB": 1 / 1000, "MB": 1, "GB": 1000, "B": 1 / 2**20}
		for unit, k in units.items():
			if value.endswith(unit):
				return float(value[: -len(unit)]) * k
		return 0.0

	def run(self):
		while not self.stop_event.is_set():
			out = subprocess.run(
				["docker", "stats", "--no-stream", "--format", "{{json .}}", self.container],
				capture_output=True, text=True,
			).stdout.strip()
			if out:
				s = json.loads(out)
				self.cpu = max(self.cpu, float(s["CPUPerc"].rstrip("%")))
				self.mem_mb = max(self.mem_mb, self._to_mb(s["MemUsage"].split("/")[0].strip()))
				self.pids = max(self.pids, int(s["PIDs"]))
			self.stop_event.wait(1)


class TaskResult:
	"""Ответ /v2 в виде, понятном classify(): status_code и json()."""

	def __init__(self, status_code: int, data: dict):
		self.status_code, self._data = status_code, data

	def json(self) -> dict:
		return self._data


async def solve_v2(client: httpx.AsyncClient, api: str, payload: dict):
	created = await client.post(f"{api}/v2/tasks", json=payload)
	if created.status_code != 202:
		return created
	url = f"{api}{created.json()['url']}"
	while True:
		await asyncio.sleep(0.5)
		task = (await client.get(url)).json()
		if task["status"] == "done":
			return TaskResult(200, task["result"])
		if task["status"] not in ("queued", "running"):
			return TaskResult(200, {"status": "err", "code": task["status"], "error": task.get("error")})


async def one(args, payload: dict) -> tuple[str, float]:
	# у каждого запроса своя сессия: общий клиент запомнил бы cookie session_id, и задачи одной
	# сессии выполнялись бы по очереди
	headers = {"X-API-Key": args.api_key} if args.api_key else {}
	t0 = time.monotonic()
	async with httpx.AsyncClient(timeout=args.http_timeout, trust_env=False, headers=headers) as client:
		try:
			if args.mode == "v2":
				resp = await solve_v2(client, args.api, payload)
			else:
				resp = await client.post(f"{args.api}/v1/solve", json=payload)
		except Exception as e:
			resp = e
	return classify(resp), time.monotonic() - t0


async def run_level(n: int, args) -> None:
	sampler = StatsSampler(args.container)
	sampler.start()
	t0 = time.monotonic()
	results = await asyncio.gather(*(one(args, build_payload(args)) for _ in range(n)))
	wall = time.monotonic() - t0
	sampler.stop_event.set()
	sampler.join()

	outcomes: dict[str, int] = {}
	for outcome, _ in results:
		outcomes[outcome] = outcomes.get(outcome, 0) + 1
	ok_times = [d for o, d in results if o == "ok"]
	median = f"{statistics.median(ok_times):.1f}" if ok_times else "-"
	print(
		f"| {n:>3} | {outcomes.get('ok', 0):>3}/{n:<3} | {median:>8} | {wall:>7.1f} "
		f"| {sampler.cpu:>7.0f} | {sampler.mem_mb / 1024:>7.2f} | {sampler.pids:>5} | "
		+ ", ".join(f"{k}={v}" for k, v in sorted(outcomes.items()) if k != "ok"),
		flush=True,
	)


async def main():
	p = argparse.ArgumentParser()
	p.add_argument("--api", default="http://127.0.0.1:8000")
	p.add_argument("--url", default="https://freelancehunt.com/projects")
	p.add_argument("--selector", default='a[href*="/project/"]')
	p.add_argument("--proxy", default="http://proxy:8080")
	p.add_argument("--captcha-type", default="cloudflare")
	p.add_argument("--mode", choices=["v1", "v2"], default="v1", help="v1 — /v1/solve, v2 — /v2/tasks с опросом")
	p.add_argument("--api-key", default=os.getenv("ROBOTEX_API_KEY"), help="X-API-Key (или ROBOTEX_API_KEY)")
	# Chrome работает в воркере — его docker stats и снимаем
	p.add_argument("--container", default="robotex-worker-1")
	p.add_argument("--levels", type=int, nargs="+", default=[1, 5, 10])
	p.add_argument("--pause", type=float, default=10, help="пауза между уровнями, с")
	p.add_argument("--http-timeout", type=float, default=300)
	args = p.parse_args()

	print("| par |   ok    | med ok,s | wall,s  | peak CPU% | peak GiB | PIDs | прочее")
	print("|-----|---------|----------|---------|-----------|----------|------|-------")
	for i, n in enumerate(args.levels):
		if i:
			await asyncio.sleep(args.pause)
		await run_level(n, args)


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