remnawave-minishop/scripts/perf_benchmarks.py

from __future__ import annotations

import argparse
import asyncio
import json
import math
import sys
import time
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

ROOT = Path(__file__).resolve().parents[1]
BACKEND = ROOT / "backend"
for path in (str(BACKEND), str(ROOT)):
    if path not in sys.path:
        sys.path.insert(0, path)

from bot.services import panel_api_service  # noqa: E402
from bot.services.tariff_worker import TariffTrafficWorker  # noqa: E402
from bot.utils import config_link  # noqa: E402
from bot.utils.config_link import prepare_config_links  # noqa: E402
from bot.utils.ttl_cache import AsyncTTLCache  # noqa: E402

DEFAULT_USER_SIZES = (200, 500, 1000, 5000, 10000)


class FakePanel:
    def __init__(self, users: int):
        self.calls = 0
        self.stats = {
            "topUsers": [
                {
                    "username": f"user_{index}",
                    "total": index + 1,
                }
                for index in range(users)
            ]
        }

    async def get_node_users_bandwidth_stats(self, node_uuid: str, *, start: str, end: str):
        self.calls += 1
        return self.stats


class FakeBulkPanel:
    def __init__(self, users: int):
        self.calls = 0
        self.users = [
            {"uuid": f"panel-{index}", "username": f"user_{index}"} for index in range(users)
        ]

    async def get_all_panel_users(self, page_size: int = 100, log_responses: bool = False):
        self.calls += 1
        return self.users


async def bench_premium_usage(users: int) -> dict:
    panel = FakePanel(users)
    worker = TariffTrafficWorker(
        settings=SimpleNamespace(),
        session_factory=SimpleNamespace(),
        panel_service=panel,
        subscription_service=SimpleNamespace(),
    )
    started = time.perf_counter()
    checksum = 0
    for index in range(users):
        checksum += await worker._premium_usage_for_user(
            f"uuid_{index}",
            ["node-1"],
            "2026-05-01",
            "2026-05-20",
            panel_username=f"user_{index}",
        )
    elapsed = time.perf_counter() - started
    return {
        "seconds": elapsed,
        "panel_calls": panel.calls,
        "checksum": checksum,
    }


async def bench_panel_user_prefetch(users: int) -> dict:
    panel = FakeBulkPanel(users)
    worker = TariffTrafficWorker(
        settings=SimpleNamespace(TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD=200),
        session_factory=SimpleNamespace(),
        panel_service=panel,
        subscription_service=SimpleNamespace(),
    )
    subs = [SimpleNamespace(panel_user_uuid=f"panel-{index}") for index in range(users)]
    started = time.perf_counter()
    by_uuid = await worker._prefetch_panel_users_by_uuid(subs)
    elapsed = time.perf_counter() - started
    return {
        "seconds": elapsed,
        "service_calls": panel.calls,
        "matched": len(by_uuid or {}),
        "legacy_user_get_calls": users,
        "estimated_bulk_http_pages_at_100": math.ceil(users / 100),
    }


async def bench_ttl_singleflight(users: int) -> dict:
    settings = SimpleNamespace(REDIS_URL="redis://example", REDIS_KEY_PREFIX="bench")
    cache = AsyncTTLCache(ttl_seconds=60, settings=settings, namespace="singleflight")
    calls = 0

    async def loader():
        nonlocal calls
        calls += 1
        await asyncio.sleep(0.001)
        return {"value": 42}

    async def fake_get(settings, key):
        return None

    async def fake_set(settings, key, value, ttl):
        return None

    started = time.perf_counter()
    with (
        patch("bot.infra.redis.cache_get_json", new=fake_get),
        patch("bot.infra.redis.cache_set_json", new=fake_set),
    ):
        await asyncio.gather(*(cache.get_or_load("same-key", loader) for _ in range(users)))
    elapsed = time.perf_counter() - started
    return {
        "seconds": elapsed,
        "loader_calls": calls,
    }


async def bench_crypt4(users: int) -> dict:
    config_link._CRYPT4_LINK_CACHES.clear()
    settings = SimpleNamespace(
        CRYPT4_ENABLED=True,
        CRYPT4_REDIRECT_URL="",
        CRYPT4_LINK_CACHE_TTL_SECONDS=3600,
        PANEL_API_URL="https://panel.example.test/api",
        PANEL_API_KEY="key",
        USER_HWID_DEVICE_LIMIT=None,
    )
    calls = 0

    async def fake_encrypt(self, raw_link: str):
        nonlocal calls
        calls += 1
        await asyncio.sleep(0.001)
        return "happ://crypt4/encrypted"

    async def fake_close(self):
        return None

    started = time.perf_counter()
    with (
        patch.object(panel_api_service.PanelApiService, "encrypt_happ_link", new=fake_encrypt),
        patch.object(panel_api_service.PanelApiService, "close_session", new=fake_close),
    ):
        await asyncio.gather(
            *(
                prepare_config_links(settings, "https://panel.example.test/sub/user")
                for _ in range(users)
            )
        )
    elapsed = time.perf_counter() - started
    return {
        "seconds": elapsed,
        "panel_calls": calls,
    }


async def run_suite(user_sizes: tuple[int, ...]) -> dict:
    results: dict[str, dict] = {}
    for users in user_sizes:
        results[str(users)] = {
            "panel_user_bulk_prefetch": await bench_panel_user_prefetch(users),
            "premium_usage_1_node": await bench_premium_usage(users),
            "ttl_cache_cold_single_key": await bench_ttl_singleflight(users),
            "crypt4_same_link": await bench_crypt4(users),
        }
    return results


def _print_table(results: dict[str, dict]) -> None:
    print(
        "users | bulk_pages_est | premium_usage_s | premium_panel_calls | "
        "ttl_loader_calls | crypt4_panel_calls"
    )
    print("-" * 104)
    for users, data in results.items():
        print(
            f"{users:>5} | "
            f"{data['panel_user_bulk_prefetch']['estimated_bulk_http_pages_at_100']:>14} | "
            f"{data['premium_usage_1_node']['seconds']:>15.6f} | "
            f"{data['premium_usage_1_node']['panel_calls']:>19} | "
            f"{data['ttl_cache_cold_single_key']['loader_calls']:>16} | "
            f"{data['crypt4_same_link']['panel_calls']:>18}"
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run bot performance microbenchmarks.")
    parser.add_argument(
        "--users",
        default=",".join(str(value) for value in DEFAULT_USER_SIZES),
        help="Comma-separated user counts. Default: 200,500,1000,5000,10000",
    )
    parser.add_argument("--json", action="store_true", help="Print JSON only.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    user_sizes = tuple(int(part.strip()) for part in args.users.split(",") if part.strip())
    results = asyncio.run(run_suite(user_sizes))
    if args.json:
        print(json.dumps({"results": results}, ensure_ascii=False))
        return
    _print_table(results)
    print()
    print(json.dumps({"results": results}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()