refactor: improve premium squads feature performance, add benchmarks
This commit is contained in:
@@ -16,6 +16,7 @@ DB_POOL_RECYCLE_SECONDS=1800 #
|
||||
REDIS_URL=redis://redis:6379/0 # Shared Redis for FSM, rate limits, cache, locks and queues
|
||||
REDIS_KEY_PREFIX=remnawave-tg-shop # Prefix for Redis keys
|
||||
WEBAPP_ME_CACHE_TTL_SECONDS=15 # Short TTL for /api/me payload cache
|
||||
WEBAPP_DEVICES_CACHE_TTL_SECONDS=5 # Short TTL for /api/devices payload cache
|
||||
WEBAPP_RATE_LIMIT_TTL_SECONDS=60 # Redis rate-limit window
|
||||
WEBAPP_RATE_LIMIT_MAX_REQUESTS=30 # Requests per window/action/user/IP
|
||||
WEBHOOK_QUEUE_NAME=webhook-events # Redis queue for heavy webhook processing
|
||||
@@ -23,6 +24,7 @@ WEBHOOK_QUEUE_CONCURRENCY=4 #
|
||||
WORKER_PANEL_SYNC_INTERVAL_SECONDS=900 # Worker panel sync interval
|
||||
TARIFF_WORKER_LOCK_TTL_SECONDS=240 # Redis lock TTL for tariff tick
|
||||
TARIFF_WORKER_TICK_SECONDS=300 # Tariff worker tick interval
|
||||
TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD=200 # Active subs threshold to bulk-fetch panel users
|
||||
|
||||
# Localization and Display
|
||||
DEFAULT_LANGUAGE="ru" # or "en"
|
||||
@@ -295,6 +297,7 @@ TRIAL_TRAFFIC_STRATEGY="NO_RESET" #
|
||||
# Connection link handling (happ crypt4)
|
||||
CRYPT4_ENABLED=False # Enable happ crypt4 encryption for subscription URLs
|
||||
CRYPT4_REDIRECT_URL= # Base redirect to wrap the connect button, e.g. https://redir.example.com?url=
|
||||
CRYPT4_LINK_CACHE_TTL_SECONDS=3600 # Cache encrypted happ links by raw subscription URL
|
||||
|
||||
# Web Server Settings (for handling webhooks)
|
||||
WEB_SERVER_HOST="0.0.0.0"
|
||||
|
||||
@@ -15,6 +15,11 @@ async def devices_route(request: web.Request) -> web.Response:
|
||||
if not db_user or db_user.is_banned:
|
||||
return _json_error(403, "access_denied", "Access denied")
|
||||
|
||||
cache_key = redis_key(settings, "cache", "webapp", "devices", user_id)
|
||||
cached = await cache_get_json(settings, cache_key)
|
||||
if isinstance(cached, dict):
|
||||
return web.json_response({"ok": True, **cached})
|
||||
|
||||
active = await subscription_service.get_active_subscription_details(session, user_id)
|
||||
panel_user_uuid = active.get("user_id") if active else None
|
||||
if not panel_user_uuid:
|
||||
@@ -32,18 +37,22 @@ async def devices_route(request: web.Request) -> web.Response:
|
||||
|
||||
devices = _normalize_devices_response(devices_response)
|
||||
max_devices = _coerce_int_or_none(active.get("max_devices")) if active else None
|
||||
return web.json_response(
|
||||
{
|
||||
"ok": True,
|
||||
"enabled": True,
|
||||
"current_devices": len(devices),
|
||||
"max_devices": max_devices,
|
||||
"max_devices_label": _format_devices_limit(max_devices),
|
||||
"devices": [
|
||||
_serialize_device(device, index) for index, device in enumerate(devices, start=1)
|
||||
],
|
||||
}
|
||||
payload = {
|
||||
"enabled": True,
|
||||
"current_devices": len(devices),
|
||||
"max_devices": max_devices,
|
||||
"max_devices_label": _format_devices_limit(max_devices),
|
||||
"devices": [
|
||||
_serialize_device(device, index) for index, device in enumerate(devices, start=1)
|
||||
],
|
||||
}
|
||||
await cache_set_json(
|
||||
settings,
|
||||
cache_key,
|
||||
payload,
|
||||
max(1, int(getattr(settings, "WEBAPP_DEVICES_CACHE_TTL_SECONDS", 5) or 5)),
|
||||
)
|
||||
return web.json_response({"ok": True, **payload})
|
||||
|
||||
|
||||
async def disconnect_device_route(request: web.Request) -> web.Response:
|
||||
@@ -103,6 +112,7 @@ async def disconnect_device_route(request: web.Request) -> web.Response:
|
||||
success = await panel_service.disconnect_device(panel_user_uuid, target_hwid)
|
||||
if not success:
|
||||
return _json_error(502, "device_disconnect_failed", "Failed to disconnect device")
|
||||
await cache_delete(settings, redis_key(settings, "cache", "webapp", "devices", user_id))
|
||||
await session.commit()
|
||||
|
||||
return web.json_response({"ok": True})
|
||||
|
||||
@@ -29,6 +29,7 @@ PREMIUM_WARNING_DEPLETED_LEVEL = PREMIUM_WARNING_LEVEL_OFFSET + 100
|
||||
# to avoid an N+1 serial chain to the Remnawave panel each tick.
|
||||
TARIFF_WORKER_BATCH_SIZE = 50
|
||||
TARIFF_WORKER_PANEL_CONCURRENCY = 10
|
||||
TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD = 200
|
||||
|
||||
|
||||
class TariffTrafficWorker:
|
||||
@@ -49,7 +50,7 @@ class TariffTrafficWorker:
|
||||
self.i18n = i18n
|
||||
self._stopped = asyncio.Event()
|
||||
self._premium_nodes_cache = {}
|
||||
self._premium_node_stats_tick_cache = {}
|
||||
self._premium_node_usage_tick_cache = {}
|
||||
|
||||
async def _user_lang(self, session: AsyncSession, user_id: int) -> str:
|
||||
try:
|
||||
@@ -133,7 +134,7 @@ class TariffTrafficWorker:
|
||||
|
||||
async def traffic_period_tick(self, session: AsyncSession) -> None:
|
||||
now = datetime.now(timezone.utc)
|
||||
self._premium_node_stats_tick_cache = {}
|
||||
self._premium_node_usage_tick_cache = {}
|
||||
warning_period_start = month_start(now)
|
||||
result = await session.execute(
|
||||
select(Subscription).where(
|
||||
@@ -146,9 +147,15 @@ class TariffTrafficWorker:
|
||||
if not subs:
|
||||
return
|
||||
|
||||
panel_users_by_uuid = await self._prefetch_panel_users_by_uuid(subs)
|
||||
semaphore = asyncio.Semaphore(TARIFF_WORKER_PANEL_CONCURRENCY)
|
||||
|
||||
async def _fetch_panel(sub: Subscription) -> dict:
|
||||
if panel_users_by_uuid is not None:
|
||||
cached_panel_user = panel_users_by_uuid.get(str(sub.panel_user_uuid))
|
||||
if cached_panel_user is not None:
|
||||
return cached_panel_user
|
||||
|
||||
async with semaphore:
|
||||
try:
|
||||
data = await self.panel_service.get_user_by_uuid(
|
||||
@@ -208,6 +215,46 @@ class TariffTrafficWorker:
|
||||
panel_user_dict=panel_data,
|
||||
)
|
||||
|
||||
async def _prefetch_panel_users_by_uuid(
|
||||
self,
|
||||
subs: list[Subscription],
|
||||
) -> Optional[dict[str, dict]]:
|
||||
threshold = int(
|
||||
getattr(
|
||||
self.settings,
|
||||
"TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD",
|
||||
TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD,
|
||||
)
|
||||
or 0
|
||||
)
|
||||
if threshold <= 0 or len(subs) < threshold:
|
||||
return None
|
||||
try:
|
||||
panel_users = await self.panel_service.get_all_panel_users(log_responses=False)
|
||||
except Exception:
|
||||
logging.exception("TariffTrafficWorker: failed to bulk-prefetch panel users")
|
||||
return None
|
||||
if not panel_users:
|
||||
return None
|
||||
|
||||
by_uuid: dict[str, dict] = {}
|
||||
for user in panel_users:
|
||||
if not isinstance(user, dict):
|
||||
continue
|
||||
uuid = user.get("uuid")
|
||||
if uuid:
|
||||
by_uuid[str(uuid)] = user
|
||||
if not by_uuid:
|
||||
return None
|
||||
matched = sum(1 for sub in subs if str(sub.panel_user_uuid) in by_uuid)
|
||||
logging.info(
|
||||
"metric panel_bulk_user_prefetch users=%s matched=%s active_subscriptions=%s",
|
||||
len(by_uuid),
|
||||
matched,
|
||||
len(subs),
|
||||
)
|
||||
return by_uuid
|
||||
|
||||
async def _ensure_period_reset_strategy(
|
||||
self,
|
||||
sub: Subscription,
|
||||
@@ -642,54 +689,95 @@ class TariffTrafficWorker:
|
||||
found = False
|
||||
username = (panel_username or "").strip() or None
|
||||
for node_uuid in node_uuids:
|
||||
stats_cache_key = (node_uuid, start_date, end_date)
|
||||
if stats_cache_key not in self._premium_node_stats_tick_cache:
|
||||
self._premium_node_stats_tick_cache[
|
||||
stats_cache_key
|
||||
] = await self.panel_service.get_node_users_bandwidth_stats(
|
||||
node_uuid,
|
||||
start=start_date,
|
||||
end=end_date,
|
||||
)
|
||||
stats = self._premium_node_stats_tick_cache.get(stats_cache_key)
|
||||
if not stats:
|
||||
lookup = await self._premium_usage_lookup_for_node(node_uuid, start_date, end_date)
|
||||
if not lookup:
|
||||
continue
|
||||
entries = stats.get("topUsers") or stats.get("usersStats") or stats.get("users") or []
|
||||
if not isinstance(entries, list):
|
||||
continue
|
||||
for entry in entries:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
user_obj = entry.get("user") if isinstance(entry.get("user"), dict) else {}
|
||||
entry_uuid = (
|
||||
user_obj.get("uuid")
|
||||
or entry.get("userUuid")
|
||||
or entry.get("uuid")
|
||||
or entry.get("user_uuid")
|
||||
|
||||
uuid_total = 0
|
||||
username_total = 0
|
||||
overlap_total = 0
|
||||
if user_uuid:
|
||||
user_uuid_str = str(user_uuid)
|
||||
uuid_total = int(lookup["by_uuid"].get(user_uuid_str, 0) or 0)
|
||||
else:
|
||||
user_uuid_str = ""
|
||||
if username:
|
||||
username_total = int(lookup["by_username"].get(username, 0) or 0)
|
||||
if user_uuid_str and username:
|
||||
overlap_total = int(
|
||||
lookup["by_uuid_username"].get((user_uuid_str, username), 0) or 0
|
||||
)
|
||||
entry_username = (
|
||||
user_obj.get("username") or entry.get("username") or entry.get("userUsername")
|
||||
)
|
||||
# Remnawave's /bandwidth-stats/nodes/{uuid}/users response
|
||||
# currently exposes only {color, username, total}; match by
|
||||
# username first, fall back to UUID if a future version
|
||||
# adds it back.
|
||||
matched = False
|
||||
if entry_uuid and entry_uuid == user_uuid:
|
||||
matched = True
|
||||
elif username and entry_username and entry_username == username:
|
||||
matched = True
|
||||
if not matched:
|
||||
continue
|
||||
value = entry.get("total")
|
||||
if value is None:
|
||||
value = int(entry.get("download", 0) or 0) + int(entry.get("upload", 0) or 0)
|
||||
total += int(value or 0)
|
||||
|
||||
node_total = uuid_total + username_total - overlap_total
|
||||
if node_total or (
|
||||
user_uuid_str in lookup["by_uuid"]
|
||||
or (username and username in lookup["by_username"])
|
||||
):
|
||||
total += node_total
|
||||
found = True
|
||||
if len(node_uuids) > 1:
|
||||
await asyncio.sleep(0.1)
|
||||
return total if found else 0
|
||||
|
||||
async def _premium_usage_lookup_for_node(
|
||||
self,
|
||||
node_uuid: str,
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
) -> Optional[dict]:
|
||||
stats_cache_key = (node_uuid, start_date, end_date)
|
||||
if stats_cache_key not in self._premium_node_usage_tick_cache:
|
||||
stats = await self.panel_service.get_node_users_bandwidth_stats(
|
||||
node_uuid,
|
||||
start=start_date,
|
||||
end=end_date,
|
||||
)
|
||||
self._premium_node_usage_tick_cache[stats_cache_key] = self._build_premium_usage_lookup(
|
||||
stats
|
||||
)
|
||||
return self._premium_node_usage_tick_cache.get(stats_cache_key)
|
||||
|
||||
@staticmethod
|
||||
def _build_premium_usage_lookup(stats: Optional[dict]) -> Optional[dict]:
|
||||
if not isinstance(stats, dict):
|
||||
return None
|
||||
entries = stats.get("topUsers") or stats.get("usersStats") or stats.get("users") or []
|
||||
if not isinstance(entries, list):
|
||||
return None
|
||||
|
||||
by_uuid: dict[str, int] = {}
|
||||
by_username: dict[str, int] = {}
|
||||
by_uuid_username: dict[tuple[str, str], int] = {}
|
||||
for entry in entries:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
user_obj = entry.get("user") if isinstance(entry.get("user"), dict) else {}
|
||||
entry_uuid = (
|
||||
user_obj.get("uuid")
|
||||
or entry.get("userUuid")
|
||||
or entry.get("uuid")
|
||||
or entry.get("user_uuid")
|
||||
)
|
||||
entry_username = (
|
||||
user_obj.get("username") or entry.get("username") or entry.get("userUsername")
|
||||
)
|
||||
value = entry.get("total")
|
||||
if value is None:
|
||||
value = int(entry.get("download", 0) or 0) + int(entry.get("upload", 0) or 0)
|
||||
total = int(value or 0)
|
||||
uuid_key = str(entry_uuid) if entry_uuid else ""
|
||||
username_key = str(entry_username) if entry_username else ""
|
||||
if uuid_key:
|
||||
by_uuid[uuid_key] = by_uuid.get(uuid_key, 0) + total
|
||||
if username_key:
|
||||
by_username[username_key] = by_username.get(username_key, 0) + total
|
||||
if uuid_key and username_key:
|
||||
pair = (uuid_key, username_key)
|
||||
by_uuid_username[pair] = by_uuid_username.get(pair, 0) + total
|
||||
return {
|
||||
"by_uuid": by_uuid,
|
||||
"by_username": by_username,
|
||||
"by_uuid_username": by_uuid_username,
|
||||
}
|
||||
|
||||
async def legacy_throttle_recovery_tick(self, session: AsyncSession) -> None:
|
||||
"""Recover subscriptions throttled by older bot versions.
|
||||
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
import hashlib
|
||||
import logging
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from bot.services.panel_api_service import PanelApiService
|
||||
from bot.utils.ttl_cache import AsyncTTLCache
|
||||
from config.settings import Settings
|
||||
|
||||
_CRYPT4_LINK_CACHES: dict[tuple[int, int], AsyncTTLCache] = {}
|
||||
|
||||
|
||||
async def _encrypt_raw_link(settings: Settings, raw_link: str) -> Optional[str]:
|
||||
"""Encrypt the raw subscription URL using the panel's happ crypt4 API."""
|
||||
@@ -14,6 +18,30 @@ async def _encrypt_raw_link(settings: Settings, raw_link: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
|
||||
def _crypt4_link_cache(settings: Settings) -> Optional[AsyncTTLCache]:
|
||||
ttl_seconds = int(getattr(settings, "CRYPT4_LINK_CACHE_TTL_SECONDS", 3600) or 0)
|
||||
if ttl_seconds <= 0:
|
||||
return None
|
||||
cache_key = (id(settings), ttl_seconds)
|
||||
cache = _CRYPT4_LINK_CACHES.get(cache_key)
|
||||
if cache is None:
|
||||
cache = AsyncTTLCache(
|
||||
ttl_seconds=ttl_seconds,
|
||||
settings=settings,
|
||||
namespace="crypt4:links",
|
||||
)
|
||||
_CRYPT4_LINK_CACHES[cache_key] = cache
|
||||
return cache
|
||||
|
||||
|
||||
async def _encrypt_raw_link_cached(settings: Settings, raw_link: str) -> Optional[str]:
|
||||
cache = _crypt4_link_cache(settings)
|
||||
if cache is None:
|
||||
return await _encrypt_raw_link(settings, raw_link)
|
||||
key = hashlib.sha256(raw_link.encode("utf-8")).hexdigest()
|
||||
return await cache.get_or_load(key, lambda: _encrypt_raw_link(settings, raw_link))
|
||||
|
||||
|
||||
async def prepare_config_links(
|
||||
settings: Settings, raw_link: Optional[str]
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
@@ -35,7 +63,7 @@ async def prepare_config_links(
|
||||
button_link = cleaned
|
||||
|
||||
if settings.CRYPT4_ENABLED:
|
||||
encrypted_payload = await _encrypt_raw_link(settings, cleaned)
|
||||
encrypted_payload = await _encrypt_raw_link_cached(settings, cleaned)
|
||||
if encrypted_payload:
|
||||
display_link = encrypted_payload
|
||||
button_link = display_link
|
||||
|
||||
@@ -37,37 +37,45 @@ class AsyncTTLCache:
|
||||
return True
|
||||
|
||||
async def get_or_load(self, key: str, loader: Callable[[], Awaitable[Any]]) -> Any:
|
||||
if self.settings is not None and self.namespace:
|
||||
try:
|
||||
from bot.infra.redis import cache_get_json, cache_set_json, redis_key
|
||||
|
||||
cache_key = redis_key(self.settings, "cache", self.namespace, key)
|
||||
cached = await cache_get_json(self.settings, cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
value = await loader()
|
||||
if self._is_cacheable(value):
|
||||
await cache_set_json(
|
||||
self.settings,
|
||||
cache_key,
|
||||
value,
|
||||
max(1, int(self.ttl_seconds)),
|
||||
)
|
||||
return value
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
cached = self.get_fresh(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
lock = self._locks.setdefault(key, asyncio.Lock())
|
||||
async with lock:
|
||||
cached = self.get_fresh(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
cache_key = None
|
||||
if self.settings is not None and self.namespace:
|
||||
try:
|
||||
from bot.infra.redis import cache_get_json, redis_key
|
||||
|
||||
cache_key = redis_key(self.settings, "cache", self.namespace, key)
|
||||
cached = await cache_get_json(self.settings, cache_key)
|
||||
if cached is not None:
|
||||
if self._is_cacheable(cached):
|
||||
self._data[key] = (time.monotonic() + self.ttl_seconds, cached)
|
||||
return cached
|
||||
except Exception:
|
||||
cache_key = None
|
||||
|
||||
value = await loader()
|
||||
if self._is_cacheable(value):
|
||||
self._data[key] = (time.monotonic() + self.ttl_seconds, value)
|
||||
if cache_key is not None:
|
||||
try:
|
||||
from bot.infra.redis import cache_set_json
|
||||
|
||||
await cache_set_json(
|
||||
self.settings,
|
||||
cache_key,
|
||||
value,
|
||||
max(1, int(self.ttl_seconds)),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return value
|
||||
|
||||
def invalidate(self, key: Optional[str] = None) -> None:
|
||||
|
||||
@@ -95,6 +95,7 @@ class Settings(BaseSettings):
|
||||
REDIS_URL: Optional[str] = Field(default=None)
|
||||
REDIS_KEY_PREFIX: str = Field(default="remnawave-tg-shop")
|
||||
WEBAPP_ME_CACHE_TTL_SECONDS: int = Field(default=15)
|
||||
WEBAPP_DEVICES_CACHE_TTL_SECONDS: int = Field(default=5)
|
||||
WEBAPP_RATE_LIMIT_TTL_SECONDS: int = Field(default=60)
|
||||
WEBAPP_RATE_LIMIT_MAX_REQUESTS: int = Field(default=30)
|
||||
WEBHOOK_QUEUE_NAME: str = Field(default="webhook-events")
|
||||
@@ -102,6 +103,7 @@ class Settings(BaseSettings):
|
||||
WORKER_PANEL_SYNC_INTERVAL_SECONDS: int = Field(default=900)
|
||||
TARIFF_WORKER_LOCK_TTL_SECONDS: int = Field(default=240)
|
||||
TARIFF_WORKER_TICK_SECONDS: int = Field(default=300)
|
||||
TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD: int = Field(default=200)
|
||||
|
||||
DEFAULT_LANGUAGE: str = Field(default="ru")
|
||||
DEFAULT_CURRENCY_SYMBOL: str = Field(default="RUB")
|
||||
@@ -269,6 +271,10 @@ class Settings(BaseSettings):
|
||||
default=None,
|
||||
description="Base redirect URL used for the connect button when crypt4 is enabled",
|
||||
)
|
||||
CRYPT4_LINK_CACHE_TTL_SECONDS: int = Field(
|
||||
default=3600,
|
||||
description="TTL for cached happ crypt4 encryption results keyed by raw subscription URL",
|
||||
)
|
||||
|
||||
WEB_SERVER_HOST: str = Field(default="0.0.0.0")
|
||||
WEB_SERVER_PORT: int = Field(default=8080)
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
"format:check:py": "python -m ruff format --check .",
|
||||
"format:check:js": "npm --prefix frontend run format:check",
|
||||
"format:check": "npm run format:check:py && npm run format:check:js",
|
||||
"bench:bot": "python scripts/perf_benchmarks.py",
|
||||
"test": "pytest -q",
|
||||
"check": "npm run lint && npm test && npm run build:webapp"
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
@@ -0,0 +1,84 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from bot.services import panel_api_service
|
||||
from bot.utils import config_link
|
||||
from bot.utils.config_link import prepare_config_links
|
||||
from bot.utils.ttl_cache import AsyncTTLCache
|
||||
|
||||
|
||||
class AsyncTTLCacheSingleflightTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_redis_backed_cache_uses_local_singleflight_on_cold_miss(self):
|
||||
settings = SimpleNamespace(REDIS_URL="redis://example", REDIS_KEY_PREFIX="test")
|
||||
cache = AsyncTTLCache(ttl_seconds=60, settings=settings, namespace="bench")
|
||||
loader_calls = 0
|
||||
set_calls = 0
|
||||
|
||||
async def loader():
|
||||
nonlocal loader_calls
|
||||
loader_calls += 1
|
||||
await asyncio.sleep(0.001)
|
||||
return {"ok": True}
|
||||
|
||||
async def fake_get(settings, key):
|
||||
return None
|
||||
|
||||
async def fake_set(settings, key, value, ttl):
|
||||
nonlocal set_calls
|
||||
set_calls += 1
|
||||
|
||||
with (
|
||||
patch("bot.infra.redis.cache_get_json", new=fake_get),
|
||||
patch("bot.infra.redis.cache_set_json", new=fake_set),
|
||||
):
|
||||
values = await asyncio.gather(*(cache.get_or_load("same", loader) for _ in range(100)))
|
||||
|
||||
self.assertEqual(values, [{"ok": True}] * 100)
|
||||
self.assertEqual(loader_calls, 1)
|
||||
self.assertEqual(set_calls, 1)
|
||||
|
||||
|
||||
class Crypt4LinkCacheTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncTearDown(self):
|
||||
config_link._CRYPT4_LINK_CACHES.clear()
|
||||
|
||||
async def test_prepare_config_links_singleflights_same_crypt4_link(self):
|
||||
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,
|
||||
)
|
||||
encrypt_calls = 0
|
||||
|
||||
async def fake_encrypt(self, raw_link):
|
||||
nonlocal encrypt_calls
|
||||
encrypt_calls += 1
|
||||
await asyncio.sleep(0.001)
|
||||
return "happ://crypt4/encrypted"
|
||||
|
||||
async def fake_close(self):
|
||||
return None
|
||||
|
||||
with (
|
||||
patch.object(panel_api_service.PanelApiService, "encrypt_happ_link", fake_encrypt),
|
||||
patch.object(panel_api_service.PanelApiService, "close_session", fake_close),
|
||||
):
|
||||
values = await asyncio.gather(
|
||||
*(
|
||||
prepare_config_links(settings, "https://panel.example.test/sub/user")
|
||||
for _ in range(100)
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(values, [("happ://crypt4/encrypted", "happ://crypt4/encrypted")] * 100)
|
||||
self.assertEqual(encrypt_calls, 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -531,3 +531,89 @@ class TariffWorkerTests(unittest.IsolatedAsyncioTestCase):
|
||||
# 4 GB used vs 1 GB baseline + 10 GB bonus = 11 GB limit → not limited.
|
||||
self.assertFalse(sub.premium_is_limited)
|
||||
self.assertEqual(int(sub.premium_used_bytes), 4 * (1024**3))
|
||||
|
||||
async def test_premium_usage_lookup_sums_uuid_and_username_without_double_counting(self):
|
||||
panel_service = AsyncMock(spec=PanelApiService)
|
||||
panel_service.get_node_users_bandwidth_stats = AsyncMock(
|
||||
return_value={
|
||||
"topUsers": [
|
||||
{"user": {"uuid": "u-1", "username": "alice"}, "total": 10},
|
||||
{"username": "alice", "total": 5},
|
||||
{"userUuid": "u-1", "total": 7},
|
||||
{"user": {"uuid": "other", "username": "alice"}, "total": 3},
|
||||
]
|
||||
}
|
||||
)
|
||||
worker = TariffTrafficWorker(
|
||||
settings=SimpleNamespace(),
|
||||
session_factory=SimpleNamespace(),
|
||||
panel_service=panel_service,
|
||||
subscription_service=SimpleNamespace(),
|
||||
)
|
||||
|
||||
total = await worker._premium_usage_for_user(
|
||||
"u-1",
|
||||
["node-1"],
|
||||
"2026-05-01",
|
||||
"2026-05-20",
|
||||
panel_username="alice",
|
||||
)
|
||||
total_again = await worker._premium_usage_for_user(
|
||||
"u-1",
|
||||
["node-1"],
|
||||
"2026-05-01",
|
||||
"2026-05-20",
|
||||
panel_username="alice",
|
||||
)
|
||||
|
||||
# The first row has both uuid and username, so it should be counted once.
|
||||
self.assertEqual(total, 25)
|
||||
self.assertEqual(total_again, 25)
|
||||
panel_service.get_node_users_bandwidth_stats.assert_awaited_once()
|
||||
|
||||
async def test_bulk_panel_prefetch_maps_panel_users_by_uuid_above_threshold(self):
|
||||
settings = SimpleNamespace(TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD=2)
|
||||
panel_service = AsyncMock(spec=PanelApiService)
|
||||
panel_service.get_all_panel_users = AsyncMock(
|
||||
return_value=[
|
||||
{"uuid": "panel-1", "username": "one"},
|
||||
{"uuid": "panel-2", "username": "two"},
|
||||
{"username": "missing-uuid"},
|
||||
]
|
||||
)
|
||||
worker = TariffTrafficWorker(
|
||||
settings=settings,
|
||||
session_factory=SimpleNamespace(),
|
||||
panel_service=panel_service,
|
||||
subscription_service=SimpleNamespace(),
|
||||
)
|
||||
|
||||
result = await worker._prefetch_panel_users_by_uuid(
|
||||
[
|
||||
SimpleNamespace(panel_user_uuid="panel-1"),
|
||||
SimpleNamespace(panel_user_uuid="panel-2"),
|
||||
]
|
||||
)
|
||||
|
||||
self.assertEqual(set(result), {"panel-1", "panel-2"})
|
||||
panel_service.get_all_panel_users.assert_awaited_once_with(log_responses=False)
|
||||
|
||||
async def test_bulk_panel_prefetch_skips_below_threshold(self):
|
||||
settings = SimpleNamespace(TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD=3)
|
||||
panel_service = AsyncMock(spec=PanelApiService)
|
||||
worker = TariffTrafficWorker(
|
||||
settings=settings,
|
||||
session_factory=SimpleNamespace(),
|
||||
panel_service=panel_service,
|
||||
subscription_service=SimpleNamespace(),
|
||||
)
|
||||
|
||||
result = await worker._prefetch_panel_users_by_uuid(
|
||||
[
|
||||
SimpleNamespace(panel_user_uuid="panel-1"),
|
||||
SimpleNamespace(panel_user_uuid="panel-2"),
|
||||
]
|
||||
)
|
||||
|
||||
self.assertIsNone(result)
|
||||
panel_service.get_all_panel_users.assert_not_awaited()
|
||||
|
||||
Reference in New Issue
Block a user