refactor: improve premium squads feature performance, add benchmarks

This commit is contained in:
3252a8
2026-05-20 22:31:37 +03:00
parent 5a0e0033ec
commit 8192eaf55b
10 changed files with 618 additions and 77 deletions
+3
View File
@@ -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"
+21 -11
View File
@@ -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})
+133 -45
View File
@@ -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.
+29 -1
View File
@@ -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
+28 -20
View File
@@ -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:
+6
View File
@@ -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)
+1
View File
@@ -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"
}
+227
View File
@@ -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()
+84
View File
@@ -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()
+86
View File
@@ -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()