refactor: improve premium squads feature performance, add benchmarks
This commit is contained in:
3252a8
@@ -0,0 +1,227 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user