diff --git a/.env.example b/.env.example index d7f53c8..93ebc23 100644 --- a/.env.example +++ b/.env.example @@ -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" diff --git a/backend/bot/app/web/webapp/devices.py b/backend/bot/app/web/webapp/devices.py index 9760f71..88f00d1 100644 --- a/backend/bot/app/web/webapp/devices.py +++ b/backend/bot/app/web/webapp/devices.py @@ -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}) diff --git a/backend/bot/services/tariff_worker.py b/backend/bot/services/tariff_worker.py index a297fae..86be457 100644 --- a/backend/bot/services/tariff_worker.py +++ b/backend/bot/services/tariff_worker.py @@ -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. diff --git a/backend/bot/utils/config_link.py b/backend/bot/utils/config_link.py index d72615d..1e1ca8b 100644 --- a/backend/bot/utils/config_link.py +++ b/backend/bot/utils/config_link.py @@ -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 diff --git a/backend/bot/utils/ttl_cache.py b/backend/bot/utils/ttl_cache.py index a65ad6a..6ba79db 100644 --- a/backend/bot/utils/ttl_cache.py +++ b/backend/bot/utils/ttl_cache.py @@ -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: diff --git a/backend/config/settings.py b/backend/config/settings.py index 5948a2e..94446db 100644 --- a/backend/config/settings.py +++ b/backend/config/settings.py @@ -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) diff --git a/package.json b/package.json index 581c0f4..80dfa3d 100644 --- a/package.json +++ b/package.json @@ -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" } diff --git a/scripts/perf_benchmarks.py b/scripts/perf_benchmarks.py new file mode 100644 index 0000000..83d51d3 --- /dev/null +++ b/scripts/perf_benchmarks.py @@ -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() diff --git a/tests/test_performance_caches.py b/tests/test_performance_caches.py new file mode 100644 index 0000000..203780b --- /dev/null +++ b/tests/test_performance_caches.py @@ -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() diff --git a/tests/test_tariff_worker.py b/tests/test_tariff_worker.py index 070e2f6..ab65078 100644 --- a/tests/test_tariff_worker.py +++ b/tests/test_tariff_worker.py @@ -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()