diff --git a/app/db/crud/user.py b/app/db/crud/user.py index 46e5def36..77d146322 100644 --- a/app/db/crud/user.py +++ b/app/db/crud/user.py @@ -1,9 +1,10 @@ -from collections.abc import Sequence +from collections.abc import AsyncIterator, Sequence +from contextlib import asynccontextmanager from copy import deepcopy from datetime import UTC, datetime, timedelta from typing import Literal -from sqlalchemy import and_, case, delete, desc, func, literal, not_, or_, select, update +from sqlalchemy import and_, bindparam, case, delete, desc, func, literal, not_, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload, selectinload, with_expression from sqlalchemy.sql import Select @@ -47,6 +48,7 @@ UserSortOption, ) from app.models.validators import MAX_ON_HOLD_EXPIRE_DURATION_SECONDS +from app.subscription.sub_update_buffer import flush_user_sub_updates, queue_user_sub_update from config import user_cleanup_settings from .general import ( @@ -65,11 +67,21 @@ tags_from_groups, ) -_USER_AGENT_MAX_LEN = UserSubscriptionUpdate.__table__.columns.user_agent.type.length or 512 -_SUBSCRIPTION_UPDATE_IP_MAX_LEN = UserSubscriptionUpdate.__table__.columns.ip.type.length or 64 _ONLINE_USERS_WINDOW = timedelta(minutes=2) +def _review_user_select_stmt(*, load_groups: bool = True) -> Select: + """Load relations needed for status review jobs without materializing usage logs.""" + return _build_user_select_stmt( + load_admin=True, + load_admin_role=True, + load_next_plan=True, + load_usage_logs=False, + load_groups=load_groups, + load_lifetime_used_traffic=True, + ) + + def _user_reset_traffic_subquery(): return ( select(func.coalesce(func.sum(UserUsageResetLogs.used_traffic_at_reset), 0)) @@ -592,19 +604,19 @@ async def remove_expired_users( async def get_active_to_expire_users(db: AsyncSession) -> list[User]: - stmt = _build_user_select_stmt().where(User.status == UserStatus.active).where(User.is_expired) + stmt = _review_user_select_stmt().where(User.status == UserStatus.active).where(User.is_expired) return list((await db.execute(stmt)).unique().scalars().all()) async def get_active_to_limited_users(db: AsyncSession) -> list[User]: - stmt = _build_user_select_stmt().where(User.status == UserStatus.active).where(User.is_limited) + stmt = _review_user_select_stmt().where(User.status == UserStatus.active).where(User.is_limited) return list((await db.execute(stmt)).unique().scalars().all()) async def get_on_hold_to_active_users(db: AsyncSession) -> list[User]: - stmt = _build_user_select_stmt().where(User.status == UserStatus.on_hold).where(User.become_online) + stmt = _review_user_select_stmt().where(User.status == UserStatus.on_hold).where(User.become_online) return list((await db.execute(stmt)).unique().scalars().all()) @@ -637,7 +649,7 @@ async def get_users_to_reset_data_usage(db: AsyncSession) -> list[User]: ) stmt = ( - _build_user_select_stmt() + _review_user_select_stmt() .outerjoin(last_reset_subq, User.id == last_reset_subq.c.user_id) .where( User.status.in_([UserStatus.active, UserStatus.limited]), @@ -666,14 +678,13 @@ async def get_usage_percentage_reached_users(db: AsyncSession, percentage: int) ) stmt = ( - _build_user_select_stmt() + _review_user_select_stmt() .options(joinedload(User.notification_reminders)) .where(User.status == UserStatus.active) .where(User.usage_percentage >= percentage) .where(not_(existing_reminder_subq)) # Only users without existing reminders ) - # All relations (admin, next_plan, usage_logs, groups) eagerly loaded via _build_user_select_stmt return list((await db.execute(stmt)).unique().scalars().all()) @@ -694,7 +705,7 @@ async def get_days_left_reached_users(db: AsyncSession, days: int) -> list[User] ) stmt = ( - _build_user_select_stmt() + _review_user_select_stmt() .options(joinedload(User.notification_reminders)) .where(User.status == UserStatus.active) .where(User.expire.isnot(None)) @@ -702,7 +713,6 @@ async def get_days_left_reached_users(db: AsyncSession, days: int) -> list[User] .where(not_(existing_reminder_subq)) # Only users without existing reminders ) - # All relations (admin, next_plan, usage_logs, groups) eagerly loaded via _build_user_select_stmt return list((await db.execute(stmt)).unique().scalars().all()) @@ -1339,28 +1349,33 @@ async def bulk_revoke_user_sub( return users +@asynccontextmanager +async def _subscription_update_read_session(db: AsyncSession) -> AsyncIterator[AsyncSession]: + """Read committed updates without reusing the caller's pre-flush snapshot.""" + await flush_user_sub_updates() + # MySQL/MariaDB REPEATABLE READ may already have a snapshot from the + # authorization/user lookup. A separate session sees the flushed rows + # without committing or rolling back any work owned by the caller. + async with AsyncSession(bind=db.bind, expire_on_commit=False) as read_db: + yield read_db + + async def user_sub_update( db: AsyncSession, user_id: int, user_agent: str, ip: str | None = None, hwid: str | None = None ) -> None: """ - Updates the user's subscription details. + Queue a subscription-update row. The request session is unused; writes are + flushed in the background so the public /sub path stays read-mostly. Args: - db (AsyncSession): Database session. + db (AsyncSession): Database session (unused; kept for call-site compatibility). user_id (int): The user id whose subscription is to be updated. user_agent (str): The user agent string. ip (str | None): The client IP address. hwid (str | None): The hardware ID of the client. """ - # Clamp to column length; some clients send very long strings (e.g. encoded configs) as User-Agent. - sanitized_user_agent = (user_agent or "")[:_USER_AGENT_MAX_LEN] - sanitized_ip = (ip or "")[:_SUBSCRIPTION_UPDATE_IP_MAX_LEN] or None - sanitized_hwid = (hwid or "")[:256] or None - agent = UserSubscriptionUpdate( - user_id=user_id, user_agent=sanitized_user_agent, ip=sanitized_ip, hwid=sanitized_hwid - ) - db.add(agent) - await db.commit() + _ = db + await queue_user_sub_update(user_id, user_agent, ip=ip, hwid=hwid) async def get_users_sub_update_list( @@ -1372,15 +1387,16 @@ async def get_users_sub_update_list( .order_by(desc(UserSubscriptionUpdate.created_at)) ) - result = await db.execute(select(func.count()).select_from(stmt.subquery())) - count = result.scalar() or 0 + async with _subscription_update_read_session(db) as read_db: + result = await read_db.execute(select(func.count()).select_from(stmt.subquery())) + count = result.scalar() or 0 - if offset: - stmt = stmt.offset(offset) - if limit: - stmt = stmt.limit(limit) + if offset: + stmt = stmt.offset(offset) + if limit: + stmt = stmt.limit(limit) - result = (await db.execute(stmt)).unique().scalars().all() + result = (await read_db.execute(stmt)).unique().scalars().all() return result, count @@ -1424,7 +1440,8 @@ async def get_users_subscription_agent_counts( stmt = stmt.where(and_(*conditions)) stmt = stmt.group_by(UserSubscriptionUpdate.user_agent) - result = await db.execute(stmt) + async with _subscription_update_read_session(db) as read_db: + result = await read_db.execute(stmt) return [(agent, count) for agent, count in result.all()] @@ -1460,7 +1477,8 @@ async def get_users_subscription_agent_stats( .order_by(trunc_expr) ) - result = await db.execute(stmt) + async with _subscription_update_read_session(db) as read_db: + result = await read_db.execute(stmt) dialect = db.bind.dialect.name rows = [] for row in result.mappings(): @@ -1825,7 +1843,11 @@ async def start_users_expire(db: AsyncSession, users: list[User]) -> list[User]: Returns: list[User]: The updated users list. """ + if not users: + return [] + now = datetime.now(UTC) + params = [] for user in users: duration = _safe_on_hold_expire_duration(user.on_hold_expire_duration) expire_time = now + timedelta(seconds=duration) if duration is not None else None @@ -1833,16 +1855,20 @@ async def start_users_expire(db: AsyncSession, users: list[User]) -> list[User]: user.on_hold_expire_duration = None user.on_hold_timeout = None user.status = UserStatus.active - stmt = ( - update(User) - .where(User.id == user.id) - .values(expire=expire_time, on_hold_expire_duration=None, on_hold_timeout=None, status=UserStatus.active) - ) - await db.execute(stmt) - + params.append({"u_id": user.id, "u_expire": expire_time}) + + await db.execute( + update(User.__table__) + .where(User.__table__.c.id == bindparam("u_id")) + .values( + expire=bindparam("u_expire"), + on_hold_expire_duration=None, + on_hold_timeout=None, + status=UserStatus.active, + ), + params, + ) await db.commit() - for user in users: - await refresh_and_load_user(db, user) return users diff --git a/app/db/models.py b/app/db/models.py index ace16b389..952d6a38b 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -283,6 +283,21 @@ def last_traffic_reset_time(self): async def inbounds(self) -> list[str]: """Returns a flat list of all included inbound tags for enabled groups.""" + loaded_groups = self.__dict__.get("groups") + if loaded_groups is not None: + inbound_tags: set[str] = set() + inbounds_loaded = True + for group in loaded_groups: + if "inbounds" not in group.__dict__: + inbounds_loaded = False + break + if group.is_disabled: + continue + for inbound in group.__dict__.get("inbounds") or []: + inbound_tags.add(inbound.tag) + if inbounds_loaded: + return list(inbound_tags) + session = async_object_session(self) if session is not None: stmt = ( @@ -299,7 +314,7 @@ async def inbounds(self) -> list[str]: # Fallback for detached instances: use already-loaded attrs only. included_tags = set() - for group in self.__dict__.get("groups") or []: + for group in loaded_groups or []: if group.is_disabled: continue for inbound in group.__dict__.get("inbounds") or []: diff --git a/app/jobs/cleanup_node_sync.py b/app/jobs/cleanup_node_sync.py new file mode 100644 index 000000000..6757401ae --- /dev/null +++ b/app/jobs/cleanup_node_sync.py @@ -0,0 +1,42 @@ +"""Compact completed node-sync work on the scheduler leader only.""" + +import asyncio + +from app import scheduler +from app.nats import needs_shared_bridge_memory +from app.nats.client import create_nats_client, get_jetstream_context +from app.nats.kv_cleanup import compact_deleted_keys +from app.nats.leader import is_job_leader +from app.utils.logger import get_logger +from config import nats_settings, runtime_settings + +logger = get_logger("jobs") + + +async def cleanup_node_sync(): + if not is_job_leader(): + return + nc = await create_nats_client() + if nc is None: + return + try: + async with asyncio.timeout(240): + count = await compact_deleted_keys(await get_jetstream_context(nc), nats_settings.node_user_sync_kv_bucket) + if count: + logger.info("Compacted %s completed node-sync keys", count) + except TimeoutError: + logger.warning("node-sync key compaction timed out; retrying on the next interval") + finally: + await nc.close() + + +if runtime_settings.role.runs_node and needs_shared_bridge_memory(): + scheduler.add_job( + cleanup_node_sync, + "interval", + seconds=300, + max_instances=1, + coalesce=True, + id="cleanup_node_sync", + replace_existing=True, + ) diff --git a/app/jobs/cleanup_subscription_updates.py b/app/jobs/cleanup_subscription_updates.py index 0865a7585..aa64432c8 100644 --- a/app/jobs/cleanup_subscription_updates.py +++ b/app/jobs/cleanup_subscription_updates.py @@ -3,6 +3,7 @@ from app import scheduler from app.db import GetDB from app.db.models import UserSubscriptionUpdate +from app.subscription.sub_update_buffer import flush_user_sub_updates from app.utils.logger import get_logger from config import job_settings, runtime_settings, subscription_env_settings @@ -12,6 +13,8 @@ async def cleanup_user_subscription_updates(): """Clean up excess user subscription updates.""" + await flush_user_sub_updates() + async with GetDB() as db: # First query: Find users that have more than the limit users_with_excess = await db.execute( diff --git a/app/jobs/record_usages.py b/app/jobs/record_usages.py index 0151fabbb..928871a50 100644 --- a/app/jobs/record_usages.py +++ b/app/jobs/record_usages.py @@ -1,9 +1,7 @@ import asyncio -import multiprocessing import random import time from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime as dt, timedelta as td from operator import attrgetter @@ -16,7 +14,7 @@ from sqlalchemy.exc import DatabaseError, OperationalError from sqlalchemy.sql.expression import Insert -from app import on_shutdown, scheduler +from app import scheduler from app.db import GetDB from app.db.base import engine from app.db.models import Admin, Node, NodeUsage, NodeUserUsage, System, User @@ -30,7 +28,8 @@ # Hard-limit concurrency: Prevent DB lock storms # Start with 2-4, adjust based on DB performance JOB_SEM = asyncio.Semaphore(3) # Max 3 concurrent DB write operations -API_SEM = asyncio.Semaphore(10) # Max 10 +API_SEM = asyncio.Semaphore(10) # Max 10 concurrent node stats RPCs +USAGE_COEFFICIENT_TTL_S = 60.0 NODE_USER_USAGE_BATCH_SIZE_BY_DIALECT = { "mysql": 1_000, "sqlite": 400, @@ -46,62 +45,7 @@ # node stats calls take longer than the scheduler interval. _user_usage_running = False _node_usage_running = False - -# Thread pool executor for I/O-bound node API calls -# Distributes workload across threads/cores for data collection -_thread_pool = None -_thread_pool_lock = asyncio.Lock() - - -async def _get_thread_pool(): - """Get or create the thread pool executor (thread-safe).""" - global _thread_pool - async with _thread_pool_lock: - if _thread_pool is None: - # Use more threads for I/O-bound operations (2x CPU cores, cap at 16) - num_workers = min(multiprocessing.cpu_count() * 2, 16) - _thread_pool = ThreadPoolExecutor(max_workers=num_workers) - logger.debug(f"Initialized ThreadPoolExecutor with {num_workers} workers") - return _thread_pool - - -@on_shutdown -async def _cleanup_thread_pool(): - """Cleanup thread pool on shutdown (thread-safe).""" - global _thread_pool - async with _thread_pool_lock: - if _thread_pool is not None: - logger.debug("Shutting down ThreadPoolExecutor...") - _thread_pool.shutdown(wait=True) - _thread_pool = None - logger.debug("ThreadPoolExecutor shut down successfully") - - -# Helper functions for threading (lightweight operations that release GIL) -def _process_node_chunk(chunk_data: tuple) -> dict: - """ - Process a chunk of node data - lightweight CPU operation. - Uses simple arithmetic and dict operations that release GIL, perfect for threads. - """ - _node_id, params, coeff = chunk_data - users_usage = defaultdict(int) - for param in params: - uid = int(param["uid"]) - value = int(param["value"] * coeff) - users_usage[uid] += value - return dict(users_usage) - - -def _merge_usage_dicts(dicts: list[dict]) -> dict: - """ - Merge multiple usage dictionaries. - Dict operations release GIL, perfect for ThreadPoolExecutor. - """ - merged = defaultdict(int) - for d in dicts: - for uid, value in d.items(): - merged[uid] += value - return dict(merged) +_usage_coefficient_cache: dict[int, tuple[float, float]] = {} def _chunked(items: list, size: int): @@ -532,29 +476,85 @@ def _process_users_stats_response(stats_response): for uid, value in params.items(): try: validated_params.append({"uid": int(uid), "value": value}) - except (ValueError, TypeError): + except ValueError, TypeError: invalid_uids.append(uid) return validated_params, invalid_uids +def _usage_job_hint(interval_env: str, interval: int) -> str: + return ( + f"Lengthen {interval_env} (currently {interval}s) or cut node stats RPC latency. " + "Raising UVICORN_WORKERS will not help — only one worker records usage." + ) + + +async def _await_usage_job(job_name: str, impl, interval: int, interval_env: str) -> None: + # No global wait_for kill: get_stats uses reset=True, so cancelling mid-run drops traffic. + start = time.monotonic() + try: + await impl() + except asyncio.CancelledError: + logger.warning("%s was cancelled", job_name) + elapsed = time.monotonic() - start + if interval > 0 and elapsed > interval: + logger.warning( + "%s took %.1fs which exceeds the %ss interval; later ticks will be skipped until this run finishes. %s", + job_name, + elapsed, + interval, + _usage_job_hint(interval_env, interval), + ) + + +async def _node_usage_coefficient(node: PasarGuardNode, node_id: int) -> float: + now = time.monotonic() + cached = _usage_coefficient_cache.get(node_id) + if cached is not None and cached[1] > now: + return cached[0] + try: + extra = await node.get_extra() + coeff = float(extra.get("usage_coefficient", 1) or 1) if extra else 1.0 + except Exception as exc: + logger.warning("Failed to get extra data for node %s: %s", node_id, exc) + coeff = cached[0] if cached is not None else 1.0 + _usage_coefficient_cache[node_id] = (coeff, now + USAGE_COEFFICIENT_TTL_S) + return coeff + + +async def _collect_node_user_usage(node: PasarGuardNode, node_id: int) -> tuple[int, float, list]: + """Fetch coefficient and user stats under one RPC slot so extra+stats overlap.""" + async with API_SEM: + coeff_result, stats_result = await asyncio.gather( + _node_usage_coefficient(node, node_id), + get_users_stats(node, node_id), + return_exceptions=True, + ) + if isinstance(coeff_result, Exception): + logger.warning("Failed to get extra data for node %s: %s", node_id, coeff_result) + coeff = 1.0 + else: + coeff = coeff_result + if isinstance(stats_result, Exception): + logger.warning("Failed to get stats for node %s: %s", node_id, stats_result) + stats: list = [] + else: + stats = stats_result + return node_id, coeff, stats + + +async def _bounded_node_rpc(coro): + async with API_SEM: + return await coro + + async def get_users_stats(node: PasarGuardNode, node_id: int | None = None): - """ - Get user stats from node using thread pool for CPU-bound processing. - This distributes the heavy data processing workload across cores. - """ + """Fetch and fold user stats from one node. Dict folding stays on the event loop.""" node_label = node_id if node_id is not None else getattr(node, "node_id", "unknown") try: - # I/O operation: fetch stats from node (async, non-blocking) - async with API_SEM: - stats_response = await node.get_stats(stat_type=StatType.UsersStat, reset=True, timeout=30) - - # CPU-bound operation: process stats in thread pool to utilize multiple cores - loop = asyncio.get_running_loop() - thread_pool = await _get_thread_pool() - validated_params, invalid_uids = await loop.run_in_executor( - thread_pool, _process_users_stats_response, stats_response - ) + # Caller holds API_SEM so extra+stats can share one slot without deadlock. + stats_response = await node.get_stats(stat_type=StatType.UsersStat, reset=True, timeout=30) + validated_params, invalid_uids = _process_users_stats_response(stats_response) if invalid_uids: for uid in invalid_uids: @@ -570,10 +570,7 @@ async def get_users_stats(node: PasarGuardNode, node_id: int | None = None): def _process_outbounds_stats_response(stats_response): - """ - Process outbounds stats response (CPU-bound operation) - can run in thread pool. - Extracted to separate function for threading. - """ + """Fold outbound uplink/downlink stats into per-row params.""" params = [ {"up": stat.value, "down": 0} if stat.type == "uplink" else {"up": 0, "down": stat.value} for stat in filter(attrgetter("value"), stats_response.stats) @@ -582,22 +579,12 @@ def _process_outbounds_stats_response(stats_response): async def get_outbounds_stats(node: PasarGuardNode, node_id: int | None = None): - """ - Get outbounds stats from node using thread pool for CPU-bound processing. - This distributes the heavy data processing workload across cores. - """ + """Fetch and fold outbound stats from one node. Dict folding stays on the event loop.""" node_label = node_id if node_id is not None else getattr(node, "node_id", "unknown") try: - # I/O operation: fetch stats from node (async, non-blocking) - async with API_SEM: - stats_response = await node.get_stats(stat_type=StatType.Outbounds, reset=True, timeout=10) - - # CPU-bound operation: process stats in thread pool to utilize multiple cores - loop = asyncio.get_running_loop() - thread_pool = await _get_thread_pool() - params = await loop.run_in_executor(thread_pool, _process_outbounds_stats_response, stats_response) - - return params + # Caller holds API_SEM so node RPCs stay bounded. + stats_response = await node.get_stats(stat_type=StatType.Outbounds, reset=True, timeout=10) + return _process_outbounds_stats_response(stats_response) except NodeAPIError as e: logger.error("Failed to get outbounds stats from node %s, error: %s", node_label, e.detail) return [] @@ -633,69 +620,19 @@ async def calculate_admin_usage(users_usage: list) -> tuple[dict, set[int]]: async def calculate_users_usage(api_params: dict, usage_coefficient: dict) -> list: - """Calculate aggregated user usage across all nodes with coefficients applied. - - Uses ThreadPoolExecutor for lightweight operations (dict/arithmetic that release GIL). - ThreadPoolExecutor is faster than ProcessPoolExecutor for these operations due to less overhead. - """ + """Aggregate user usage across nodes with coefficients applied.""" if not api_params: return [] - def _process_usage_sync(chunks_data: list[tuple[int, list[dict], float]]): - """Synchronous fallback used for small batches or on executor failures.""" - users_usage = defaultdict(int) - for _, params, coeff in chunks_data: - for param in params: - uid = int(param["uid"]) - value = int(param["value"] * coeff) - users_usage[uid] += value - return [{"uid": uid, "value": value} for uid, value in users_usage.items()] - - # Prepare chunks for parallel processing - chunks = [ - (node_id, params, usage_coefficient.get(node_id, 1)) - for node_id, params in api_params.items() - if params # Skip empty params - ] - - if not chunks: - return [] - - # For small datasets, process synchronously to avoid overhead - total_params = sum(len(params) for _, params, _ in chunks) - if total_params < 1000: - return _process_usage_sync(chunks) - - # Large dataset - use ThreadPoolExecutor (faster for lightweight operations) - loop = asyncio.get_running_loop() - try: - thread_pool = await _get_thread_pool() - except Exception: - logger.exception("Falling back to synchronous user usage calculation: failed to init thread pool") - return _process_usage_sync(chunks) + users_usage: dict[int, int] = defaultdict(int) + for node_id, params in api_params.items(): + if not params: + continue + coeff = usage_coefficient.get(node_id, 1) + for param in params: + users_usage[int(param["uid"])] += int(param["value"] * coeff) - try: - # Process chunks in parallel using threads (less overhead than processes) - tasks = [loop.run_in_executor(thread_pool, _process_node_chunk, chunk) for chunk in chunks] - chunk_results = await asyncio.gather(*tasks) - - # Merge results - also lightweight, use threads - if len(chunk_results) > 4: - # Split merge operation into smaller chunks - chunk_size = max(1, len(chunk_results) // 4) - merge_chunks = [chunk_results[i : i + chunk_size] for i in range(0, len(chunk_results), chunk_size)] - merge_tasks = [ - loop.run_in_executor(thread_pool, _merge_usage_dicts, merge_chunk) for merge_chunk in merge_chunks - ] - partial_results = await asyncio.gather(*merge_tasks) - final_result = _merge_usage_dicts(partial_results) - else: - final_result = _merge_usage_dicts(chunk_results) - - return [{"uid": uid, "value": value} for uid, value in final_result.items()] - except Exception: - logger.exception("Falling back to synchronous user usage calculation: executor merge failed") - return _process_usage_sync(chunks) + return [{"uid": uid, "value": value} for uid, value in users_usage.items()] async def _record_user_usages_impl(): @@ -713,29 +650,22 @@ async def _record_user_usages_impl(): logger.debug(f"Starting user usage recording for {len(nodes)} nodes") try: - # Gather node extra data directly without unnecessary task creation - node_data = await asyncio.gather(*[node.get_extra() for _, node in nodes], return_exceptions=True) - usage_coefficient = {} - for (node_id, _), data in zip(nodes, node_data): - if isinstance(data, Exception): - logger.warning(f"Failed to get extra data for node {node_id}: {data}") - usage_coefficient[node_id] = 1.0 - else: - usage_coefficient[node_id] = data.get("usage_coefficient", 1) if data else 1.0 - - # Gather stats directly - asyncio.gather accepts coroutines, no need for create_task - stats_results = await asyncio.gather( - *[get_users_stats(node, node_id) for node_id, node in nodes], + collected = await asyncio.gather( + *[_collect_node_user_usage(node, node_id) for node_id, node in nodes], return_exceptions=True, ) + usage_coefficient = {} api_params = {} - for i, result in enumerate(stats_results): + for i, result in enumerate(collected): node_id = nodes[i][0] if isinstance(result, Exception): - logger.warning(f"Failed to get stats for node {node_id}: {result}") + logger.warning("Failed to collect usage for node %s: %s", node_id, result) + usage_coefficient[node_id] = 1.0 api_params[node_id] = [] - else: - api_params[node_id] = result + continue + _, coeff, stats = result + usage_coefficient[node_id] = coeff + api_params[node_id] = stats users_usage = await calculate_users_usage(api_params, usage_coefficient) if not users_usage: @@ -814,22 +744,29 @@ async def _record_user_usages_impl(): async def record_user_usages(): - """ - Record user usages with hard timeout. - Jobs running longer than 2 minutes are forcefully cancelled. + """Record user usages. Overlapping ticks are skipped; there is no global kill. + + ``get_stats(..., reset=True)`` zeros node counters, so a 120s cancel after + that drop can lose traffic. If this job is skipped, lengthen + JOB_RECORD_USER_USAGES_INTERVAL or cut node RPC latency — extra Uvicorn + workers will not help. """ global _user_usage_running if _user_usage_running: - logger.warning("record_user_usages skipped; previous run still in progress") + logger.warning( + "record_user_usages skipped; previous run still in progress. %s", + _usage_job_hint("JOB_RECORD_USER_USAGES_INTERVAL", job_settings.record_user_usages_interval), + ) return _user_usage_running = True try: - await asyncio.wait_for(_record_user_usages_impl(), timeout=120) - except TimeoutError: - logger.warning("record_user_usages killed after 120s timeout") - except asyncio.CancelledError: - logger.warning("record_user_usages was cancelled") + await _await_usage_job( + "record_user_usages", + _record_user_usages_impl, + job_settings.record_user_usages_interval, + "JOB_RECORD_USER_USAGES_INTERVAL", + ) finally: _user_usage_running = False @@ -851,7 +788,7 @@ async def _record_node_usages_impl(): try: # Get healthy nodes and gather stats directly stats_results = await asyncio.gather( - *[get_outbounds_stats(node, node_id) for node_id, node in nodes], + *[_bounded_node_rpc(get_outbounds_stats(node, node_id)) for node_id, node in nodes], return_exceptions=True, ) api_params = {} @@ -924,22 +861,23 @@ async def _record_node_usages_impl(): async def record_node_usages(): - """ - Record node usages with hard timeout. - Jobs running longer than 2 minutes are forcefully cancelled. - """ + """Record node usages. Same skip rules as ``record_user_usages``.""" global _node_usage_running if _node_usage_running: - logger.warning("record_node_usages skipped; previous run still in progress") + logger.warning( + "record_node_usages skipped; previous run still in progress. %s", + _usage_job_hint("JOB_RECORD_NODE_USAGES_INTERVAL", job_settings.record_node_usages_interval), + ) return _node_usage_running = True try: - await asyncio.wait_for(_record_node_usages_impl(), timeout=120) - except TimeoutError: - logger.warning("record_node_usages killed after 120s timeout") - except asyncio.CancelledError: - logger.warning("record_node_usages was cancelled") + await _await_usage_job( + "record_node_usages", + _record_node_usages_impl, + job_settings.record_node_usages_interval, + "JOB_RECORD_NODE_USAGES_INTERVAL", + ) finally: _node_usage_running = False diff --git a/app/jobs/reset_user_data_usage.py b/app/jobs/reset_user_data_usage.py index 28e1998c2..edbbcc1ad 100644 --- a/app/jobs/reset_user_data_usage.py +++ b/app/jobs/reset_user_data_usage.py @@ -5,6 +5,7 @@ from app.db import GetDB from app.db.crud.user import bulk_reset_user_data_usage, get_users_to_reset_data_usage from app.jobs.dependencies import SYSTEM_ADMIN +from app.node.sync import sync_users from app.operation import OperatorType from app.operation.user import UserOperation from app.utils.logger import get_logger @@ -25,8 +26,9 @@ async def reset_data_usage(): clean_chart_data=usage_settings.reset_user_usage_clean_chart_data, ) + await sync_users(updated_users) for db_user in updated_users: - user = await user_operator.update_user(db_user) + user = await user_operator.validate_user(db_user) asyncio.create_task(notification.reset_user_data_usage(user, SYSTEM_ADMIN)) if old_statuses.get(user.id) != user.status: diff --git a/app/jobs/review_users.py b/app/jobs/review_users.py index cb8bacbd0..67e270e1e 100644 --- a/app/jobs/review_users.py +++ b/app/jobs/review_users.py @@ -20,6 +20,7 @@ from app.jobs.dependencies import SYSTEM_ADMIN from app.models.settings import Webhook from app.models.user import UserNotificationResponse +from app.node.sync import sync_users from app.operation import OperatorType from app.operation.user import UserOperation from app.settings import webhook_settings @@ -30,48 +31,72 @@ user_operator = UserOperation(operator_type=OperatorType.SYSTEM) -async def change_status(db: AsyncSession, db_user: User, status: UserStatus): - next_plan_activated = bool(db_user.next_plan) and status != UserStatus.active - if next_plan_activated: - db_user = await reset_user_by_next( - db, - db_user, - clean_chart_data=usage_settings.reset_user_usage_clean_chart_data, - ) +async def _notify_status_change(db_user: User, status: UserStatus) -> None: + user = await user_operator.validate_user(db_user) + asyncio.create_task(notification.user_status_change(user, SYSTEM_ADMIN)) + logger.info(f'User "{user.username}" status changed to {status.value}') + + +async def _notify_next_plan(db_user: User) -> None: + user = await user_operator.validate_user(db_user) + asyncio.create_task(notification.user_data_reset_by_next(user, SYSTEM_ADMIN)) + logger.info(f'User "{db_user.username}" next plan activated') - user = await user_operator.update_user(db_user) - if next_plan_activated: - asyncio.create_task(notification.user_data_reset_by_next(user, SYSTEM_ADMIN)) - logger.info(f'User "{db_user.username}" next plan activated') +async def apply_status_changes(db: AsyncSession, users: list[User], status: UserStatus) -> None: + """Bulk-sync status changes, only walking next-plan users one by one.""" + if not users: return - asyncio.create_task(notification.user_status_change(user, SYSTEM_ADMIN)) - logger.info(f'User "{user.username}" status changed to {status.value}') + next_plan_users: list[User] = [] + plain_users: list[User] = [] + for user in users: + if user.next_plan is not None and status != UserStatus.active: + next_plan_users.append(user) + else: + plain_users.append(user) + + if plain_users: + if status in (UserStatus.expired, UserStatus.limited): + await update_users_status(db, plain_users, status) + await sync_users(plain_users) + for db_user in plain_users: + await _notify_status_change(db_user, status) + + if not next_plan_users: + return + + reset_users: list[User] = [] + for db_user in next_plan_users: + reset_users.append( + await reset_user_by_next( + db, + db_user, + clean_chart_data=usage_settings.reset_user_usage_clean_chart_data, + ) + ) + await sync_users(reset_users) + for db_user in reset_users: + await _notify_next_plan(db_user) async def expire_users_job(): async with GetDB() as db: if expired_users := await get_active_to_expire_users(db): - updated_users = await update_users_status(db, expired_users, UserStatus.expired) - for user in updated_users: - await change_status(db, user, UserStatus.expired) + await apply_status_changes(db, expired_users, UserStatus.expired) async def limit_users_job(): async with GetDB() as db: if limited_users := await get_active_to_limited_users(db): - updated_users = await update_users_status(db, limited_users, UserStatus.limited) - for user in updated_users: - await change_status(db, user, UserStatus.limited) + await apply_status_changes(db, limited_users, UserStatus.limited) async def on_hold_to_active_users_job(): async with GetDB() as db: if on_hold_users := await get_on_hold_to_active_users(db): updated_users = await start_users_expire(db, on_hold_users) - for user in updated_users: - await change_status(db, user, UserStatus.active) + await apply_status_changes(db, updated_users, UserStatus.active) async def usage_percent_notification_job(): diff --git a/app/nats/kv_cas.py b/app/nats/kv_cas.py index c9c2c90e4..8e2a10210 100644 --- a/app/nats/kv_cas.py +++ b/app/nats/kv_cas.py @@ -12,6 +12,7 @@ import nats.js.errors as nats_js_errors from nats.js.kv import KeyValue +from app.nats.kv_watch import watch_kv from app.utils.logger import get_logger logger = get_logger("nats-kv-cas") @@ -70,12 +71,17 @@ async def kv_cas_json(kv: CasKv, key: str, value: dict[str, Any], revision: int) return False -async def kv_put_json(kv: CasKv, key: str, value: dict[str, Any]) -> None: +async def kv_put_json(kv: CasKv, key: str, value: dict[str, Any]) -> int: """Upsert JSON with CAS retries (latest value wins).""" + payload = json.dumps(value, separators=(",", ":")).encode() for attempt in range(32): _, rev = await kv_get_json(kv, key) - if await kv_cas_json(kv, key, value, rev): - return + try: + if rev == 0: + return await kv.create(key, payload) + return await kv.update(key, payload, last=rev) + except nats_errors.Error as exc: + logger.debug("NATS KV put attempt failed for key=%s revision=%s: %s", key, rev, exc) if attempt < 31: await cas_retry_backoff() raise RuntimeError(f"failed to put NATS KV key={key} after CAS retries") @@ -91,7 +97,7 @@ async def kv_list_keys(kv: CasKv, prefix: str) -> list[str]: # filter subject, so use that to filter server-side to this prefix only. watch = getattr(kv, "watch", None) if callable(watch): - watcher = await watch(f"{prefix}*", ignore_deletes=True, meta_only=True) + watcher = await watch_kv(kv, f"{prefix}*", ignore_deletes=True, inactive_threshold=5, snapshot_only=True) try: keys: list[str] = [] async for entry in watcher: diff --git a/app/nats/kv_cleanup.py b/app/nats/kv_cleanup.py new file mode 100644 index 000000000..1359546d2 --- /dev/null +++ b/app/nats/kv_cleanup.py @@ -0,0 +1,35 @@ +"""Remove old queue tombstones without deleting concurrently recreated keys.""" + +from __future__ import annotations + +import contextlib +from datetime import UTC, datetime + +from nats.js.client import JetStreamContext +from nats.js.errors import BucketNotFoundError + +from app.nats.kv_watch import watch_kv + + +async def compact_deleted_keys(js: JetStreamContext, bucket: str, *, older_than: float = 300) -> int: + try: + kv = await js.key_value(bucket) + except BucketNotFoundError: + return 0 + watcher = await watch_kv(kv, ">", inactive_threshold=5, snapshot_only=True) + cutoff = datetime.now(UTC).timestamp() - older_than + purged = 0 + try: + async for entry in watcher: + if entry is None: + break + if entry.operation not in ("DEL", "PURGE") or entry.created.timestamp() > cutoff: + continue + # A key can be recreated after this snapshot. Purge ONLY revisions + # up to the observed tombstone, never the newer pending update. + await js.purge_stream(f"KV_{bucket}", subject=f"$KV.{bucket}.{entry.key}", seq=entry.revision + 1) + purged += 1 + finally: + with contextlib.suppress(Exception): + await watcher.stop() + return purged diff --git a/app/nats/kv_index.py b/app/nats/kv_index.py new file mode 100644 index 000000000..4bafe528c --- /dev/null +++ b/app/nats/kv_index.py @@ -0,0 +1,140 @@ +"""A live key-only index for the shared node sync queue. + +KV values and CAS revisions still come from JetStream. The index only discovers +candidate keys, so a delayed notification cannot grant ownership of a user. +""" + +from __future__ import annotations + +import asyncio +import contextlib + +from app.nats.kv_cas import CasKv, kv_list_keys +from app.nats.kv_watch import watch_kv +from app.utils.logger import get_logger + +logger = get_logger("nats-kv-index") + + +class KvKeyIndex: + SNAPSHOT_STALL_TIMEOUT = 30 + + def __init__(self, kv: CasKv): + self._kv = kv + self._keys: dict[str, dict[str, int]] = {} + self._local_puts: dict[str, int] = {} + self._last_revision = 0 + self._ready = asyncio.Event() + self._task: asyncio.Task | None = None + self._closed = False + + async def entries(self, prefix: str) -> dict[str, int]: + if self._closed: + raise RuntimeError("NATS key index is closed") + if not callable(getattr(self._kv, "watch", None)): + return dict.fromkeys(await kv_list_keys(self._kv, prefix), 0) + if self._task is None: + # No await between checking and setting: all node callers share one + # watcher per process, including concurrent first-time callers. + self._task = asyncio.create_task(self._run(), name="node-sync-key-index") + while not self._ready.is_set(): + revision = self._last_revision + try: + async with asyncio.timeout(self.SNAPSHOT_STALL_TIMEOUT): + await self._ready.wait() + except TimeoutError: + # Large snapshots can take longer than one timeout interval. + # Fail only when replay has stopped making progress. + if self._last_revision <= revision: + raise + if self._closed: + raise RuntimeError("NATS key index is closed") + return dict(self._keys.get(prefix, {})) + + async def keys(self, prefix: str) -> list[str]: + return list(await self.entries(prefix)) + + async def snapshot(self, prefixes: tuple[str, ...]) -> tuple[set[str], int]: + """Copy candidate keys and their replay checkpoint without yielding.""" + await self.entries(prefixes[0]) + return {key for prefix in prefixes for key in self._keys.get(prefix, {})}, self._last_revision + + def observe_put(self, key: str, revision: int) -> None: + """Expose an acknowledged local write before its watch event arrives. + + The bridge can exit its lazy sync loop on an empty claim, so local + enqueue/requeue must be immediately discoverable. Ignore old watch + events until they catch up with this write's revision. + """ + if ( + self._closed + or not callable(getattr(self._kv, "watch", None)) + or revision <= max(self._last_revision, self._local_puts.get(key, 0)) + ): + return + self._local_puts[key] = revision + self._keys.setdefault(key.rpartition(".")[0] + ".", {})[key] = revision + + def discard(self, key: str, revision: int) -> None: + """Evict a stale candidate only if a newer event has not replaced it.""" + prefix = key.rpartition(".")[0] + "." + keys = self._keys.get(prefix) + if keys is not None and keys.get(key) == revision: + del keys[key] + if not keys: + del self._keys[prefix] + + async def _run(self) -> None: + retry_delay = 0.5 + while not self._closed: + watcher = None + self._ready.clear() + self._keys.clear() + self._local_puts.clear() + self._last_revision = 0 + try: + # Replay once, then consume changes. Do not ignore deletes: + # removing them from the index bounds memory by live work. + watcher = await watch_kv(self._kv, ">", inactive_threshold=30) + async for entry in watcher: + if entry is None: + self._ready.set() + retry_delay = 0.5 + continue + self._last_revision = max(self._last_revision, entry.revision) + if self._local_puts.get(entry.key, 0) > entry.revision: + continue + self._local_puts.pop(entry.key, None) + prefix, _, _ = entry.key.rpartition(".") + prefix += "." + if entry.operation in ("DEL", "PURGE"): + keys = self._keys.get(prefix) + if keys is not None: + keys.pop(entry.key, None) + if not keys: + del self._keys[prefix] + else: + self._keys.setdefault(prefix, {})[entry.key] = entry.revision + if not self._closed: + raise RuntimeError("NATS key watcher stopped") + except asyncio.CancelledError: + raise + except Exception: + logger.warning("Node sync key index interrupted; rebuilding after backoff", exc_info=True) + finally: + self._ready.clear() + if watcher is not None: + with contextlib.suppress(Exception): + await watcher.stop() + await asyncio.sleep(retry_delay) + retry_delay = min(retry_delay * 2, 30) + + async def close(self) -> None: + self._closed = True + if self._task is not None: + self._task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._task + self._keys.clear() + self._local_puts.clear() + self._ready.set() diff --git a/app/nats/kv_watch.py b/app/nats/kv_watch.py new file mode 100644 index 000000000..3ca92624f --- /dev/null +++ b/app/nats/kv_watch.py @@ -0,0 +1,159 @@ +"""Bounded KV snapshots followed by ordered live updates.""" + +from __future__ import annotations + +import asyncio +import contextlib + +from nats.js.api import AckPolicy, ConsumerConfig, DeliverPolicy +from nats.js.kv import KV_DEL, KV_MARKER_REASON, KV_OP, KV_PURGE, KeyValue + + +async def watch_kv( + kv, pattern: str, *, ignore_deletes=False, inactive_threshold=30, snapshot_only=False, start_revision=None +): + if not isinstance(kv, KeyValue): + return await kv.watch( + pattern, ignore_deletes=ignore_deletes, meta_only=True, inactive_threshold=inactive_threshold + ) + return KvWatcher(kv, pattern, ignore_deletes, inactive_threshold, snapshot_only, start_revision) + + +class KvWatcher: + """Pull snapshots in batches so slow readers cannot overflow subscriptions. + + A live ordered subscription resumes after the snapshot's last revision. + Snapshot consumers are deleted immediately, including on cancellation. + """ + + BATCH_SIZE = 256 + + def __init__(self, kv, pattern, ignore_deletes, inactive_threshold, snapshot_only, start_revision): + self._kv = kv + self._pattern = pattern + self._ignore_deletes = ignore_deletes + self._inactive_threshold = inactive_threshold + self._snapshot_only = snapshot_only + self._start_revision = start_revision + # The KV API does not expose a bounded snapshot operation. Keep its + # underlying JetStream connection and subject access in this adapter. + self._js = kv._js + self._stream = kv._stream + self._subject = kv._pre + pattern + self._snapshot = None + self._live = None + self._updates = asyncio.Queue(maxsize=self.BATCH_SIZE) + self._iterator = self._iterate() + + def __aiter__(self): + return self + + async def __anext__(self): + return await self._iterator.__anext__() + + def _entry(self, msg): + operation = None + if msg.headers: + operation = msg.headers.get(KV_OP) + if operation is None and KV_MARKER_REASON in msg.headers: + reason = msg.headers[KV_MARKER_REASON] + if reason in ("MaxAge", "Purge"): + operation = KV_PURGE + elif reason == "Remove": + operation = KV_DEL + else: + return None + if self._ignore_deletes and operation in (KV_DEL, KV_PURGE): + return None + meta = msg.metadata + return KeyValue.Entry( + bucket=self._kv._name, + key=msg.subject[len(self._kv._pre) :], + value=b"", + revision=meta.sequence.stream, + delta=meta.num_pending, + created=meta.timestamp, + operation=operation, + ) + + async def _close_subscription(self, subscription): + if subscription is not None: + with contextlib.suppress(Exception): + await subscription.unsubscribe() + with contextlib.suppress(Exception): + await self._js.delete_consumer(self._stream, subscription._consumer) + + async def _iterate(self): + nc = self._js._nc + reconnects = nc.stats["reconnects"] + try: + watermark = 0 if self._snapshot_only else (await self._js.stream_info(self._stream)).state.last_seq + self._snapshot = await self._js.pull_subscribe( + self._subject, + stream=self._stream, + config=ConsumerConfig( + deliver_policy=( + DeliverPolicy.BY_START_SEQUENCE + if self._start_revision is not None + else DeliverPolicy.LAST_PER_SUBJECT + ), + opt_start_seq=self._start_revision, + ack_policy=AckPolicy.NONE, + headers_only=True, + inactive_threshold=max(self._inactive_threshold, 60), + ), + pending_msgs_limit=self.BATCH_SIZE * 2, + pending_bytes_limit=1024 * 1024, + ) + pending = (await self._snapshot.consumer_info()).num_pending + while pending: + try: + messages = await self._snapshot.fetch(min(self.BATCH_SIZE, pending), timeout=5) + except TimeoutError: + messages = [] + # An interrupted snapshot must be rebuilt: an AckNone consumer + # may have advanced while its client was disconnected. + if not nc.is_connected or nc.stats["reconnects"] != reconnects: + raise RuntimeError("NATS connection changed during KV snapshot") + for msg in messages: + if not self._snapshot_only: + watermark = max(watermark, msg.metadata.sequence.stream) + entry = self._entry(msg) + if entry is not None: + yield entry + if messages: + pending = messages[-1].metadata.num_pending + else: + pending = (await self._snapshot.consumer_info()).num_pending + if not nc.is_connected or nc.stats["reconnects"] != reconnects: + raise RuntimeError("NATS connection changed during KV snapshot") + await self._close_subscription(self._snapshot) + self._snapshot = None + if self._snapshot_only: + yield None + return + + async def receive(msg): + entry = self._entry(msg) + if entry is not None: + await self._updates.put(entry) + + self._live = await self._js.subscribe( + self._subject, + stream=self._stream, + config=ConsumerConfig(opt_start_seq=watermark + 1), + deliver_policy=DeliverPolicy.BY_START_SEQUENCE, + ordered_consumer=True, + headers_only=True, + inactive_threshold=self._inactive_threshold, + cb=receive, + ) + yield None + while True: + yield await self._updates.get() + finally: + await self._close_subscription(self._snapshot) + await self._close_subscription(self._live) + + async def stop(self): + await self._iterator.aclose() diff --git a/app/node/nats_memory.py b/app/node/nats_memory.py index 2523344ee..47ec366a3 100644 --- a/app/node/nats_memory.py +++ b/app/node/nats_memory.py @@ -9,6 +9,7 @@ import json import os import time +from itertools import islice from typing import Any from uuid import uuid4 @@ -27,6 +28,8 @@ from app.nats import needs_shared_bridge_memory from app.nats.client import create_nats_client, get_jetstream_context, get_or_create_kv_bucket from app.nats.kv_cas import CasKv, cas_retry_backoff, kv_cas_json, kv_get_json, kv_list_keys, kv_put_json +from app.nats.kv_index import KvKeyIndex +from app.nats.kv_watch import watch_kv from app.utils.logger import get_logger from config import nats_settings @@ -96,6 +99,24 @@ class NatsUserSyncStore: def __init__(self, kv: CasKv): self._kv = kv + self._key_index = KvKeyIndex(kv) + self._write_slots = asyncio.Semaphore(32) + + async def close(self) -> None: + await self._key_index.close() + + async def _run_bounded(self, operation, items) -> None: + async def run(item): + async with self._write_slots: + await operation(item) + + items = iter(items) + while batch := list(islice(items, 32)): + # Bound both live tasks and concurrent KV operations during full syncs. + results = await asyncio.gather(*(run(item) for item in batch), return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + raise result def _pending_prefix(self, node_id: str) -> str: return f"p.{node_id}." @@ -120,20 +141,22 @@ async def _enqueue_one(self, node_id: str, email: str, user: User) -> None: key = self._pending_key(node_id, email) value = {"email": email, "user": _b64_user(user)} self._ensure_value_size(key, value) - await kv_put_json(self._kv, key, value) + revision = await kv_put_json(self._kv, key, value) + self._key_index.observe_put(key, revision) async def enqueue_users(self, node_id: str, users: list[User]) -> None: if not users: return # Latest payload per email wins (dedupe across the batch first). by_email = {user.email: user for user in users} - await asyncio.gather(*(self._enqueue_one(node_id, email, user) for email, user in by_email.items())) + await self._run_bounded(lambda user: self._enqueue_one(node_id, user.email, user), by_email.values()) async def _requeue_expired_claims(self, node_id: str) -> None: now = time.time() - for key in await kv_list_keys(self._kv, self._claimed_prefix(node_id)): + for key, indexed_revision in (await self._key_index.entries(self._claimed_prefix(node_id))).items(): doc, rev = await kv_get_json(self._kv, key) if doc is None: + self._key_index.discard(key, indexed_revision) continue if float(doc.get("expires_at", 0)) > now: continue @@ -142,7 +165,8 @@ async def _requeue_expired_claims(self, node_id: str) -> None: if isinstance(email, str) and isinstance(user_b64, str): pending_key = self._pending_key(node_id, email) pending_value = {"email": email, "user": user_b64} - await kv_put_json(self._kv, pending_key, pending_value) + revision = await kv_put_json(self._kv, pending_key, pending_value) + self._key_index.observe_put(pending_key, revision) try: await self._kv.delete(key, last=rev) except Exception as exc: @@ -155,11 +179,12 @@ async def claim_users(self, node_id: str, worker_id: str, limit: int, lease_seco result: list[ClaimedUser] = [] now = time.time() - for pending_key in await kv_list_keys(self._kv, self._pending_prefix(node_id)): + for pending_key, indexed_revision in (await self._key_index.entries(self._pending_prefix(node_id))).items(): if len(result) >= limit: break doc, rev = await kv_get_json(self._kv, pending_key) if doc is None: + self._key_index.discard(pending_key, indexed_revision) continue email = doc.get("email") user_b64 = doc.get("user") @@ -177,8 +202,14 @@ async def claim_users(self, node_id: str, worker_id: str, limit: int, lease_seco self._ensure_value_size(claimed_key, claimed_value) try: # Create-only so two workers cannot claim into the same token key. - if not await kv_cas_json(self._kv, claimed_key, claimed_value, 0): + try: + claimed_revision = await self._kv.create( + claimed_key, json.dumps(claimed_value, separators=(",", ":")).encode() + ) + except nats.errors.Error as exc: + logger.debug("Failed to create claim key=%s: %s", claimed_key, exc) continue + self._key_index.observe_put(claimed_key, claimed_revision) await self._kv.delete(pending_key, last=rev) except Exception as exc: logger.debug("Claim race for pending key=%s: %s", pending_key, exc) @@ -207,13 +238,14 @@ async def _ack_one(self, node_id: str, token: str) -> None: async def ack_users(self, node_id: str, tokens: list[str]) -> None: if not tokens: return - await asyncio.gather(*(self._ack_one(node_id, token) for token in tokens)) + await self._run_bounded(lambda token: self._ack_one(node_id, token), tokens) async def _requeue_one(self, node_id: str, item: ClaimedUser) -> None: pending_key = self._pending_key(node_id, item.user.email) pending_value = {"email": item.user.email, "user": _b64_user(item.user)} self._ensure_value_size(pending_key, pending_value) - await kv_put_json(self._kv, pending_key, pending_value) + revision = await kv_put_json(self._kv, pending_key, pending_value) + self._key_index.observe_put(pending_key, revision) claimed_key = self._claimed_key(node_id, item.token) doc, rev = await kv_get_json(self._kv, claimed_key) if doc is None: @@ -226,7 +258,7 @@ async def _requeue_one(self, node_id: str, item: ClaimedUser) -> None: async def requeue_users(self, node_id: str, claimed_users: list[ClaimedUser]) -> None: if not claimed_users: return - await asyncio.gather(*(self._requeue_one(node_id, item) for item in claimed_users)) + await self._run_bounded(lambda item: self._requeue_one(node_id, item), claimed_users) async def _clear_one(self, key: str) -> None: doc, rev = await kv_get_json(self._kv, key) @@ -238,9 +270,26 @@ async def _clear_one(self, key: str) -> None: logger.debug("Failed to clear key=%s: %s", key, exc) async def clear(self, node_id: str) -> None: - pending_keys = await kv_list_keys(self._kv, self._pending_prefix(node_id)) - claimed_keys = await kv_list_keys(self._kv, self._claimed_prefix(node_id)) - await asyncio.gather(*(self._clear_one(key) for key in [*pending_keys, *claimed_keys])) + prefixes = (self._pending_prefix(node_id), self._claimed_prefix(node_id)) + if isinstance(self._kv, KeyValue): + keys, revision = await self._key_index.snapshot(prefixes) + # Include acknowledged writes from other workers even if their live + # notifications are delayed. Replaying only the tail avoids reading + # the node's entire deletion history on every reconnect. + watcher = await watch_kv( + self._kv, f"*.{node_id}.*", ignore_deletes=True, snapshot_only=True, start_revision=revision + 1 + ) + try: + async for entry in watcher: + if entry is None: + break + if entry.key.startswith(prefixes): + keys.add(entry.key) + finally: + await watcher.stop() + else: + keys = set(await kv_list_keys(self._kv, prefixes[0])) | set(await kv_list_keys(self._kv, prefixes[1])) + await self._run_bounded(self._clear_one, keys) class NatsNodeLifecycleCoordinator: @@ -445,6 +494,8 @@ def get_bridge_memory() -> tuple[NatsUserSyncStore | None, NatsNodeLifecycleCoor async def shutdown_bridge_memory() -> None: global _nc, _user_sync_kv, _lifecycle_kv, _user_sync_store, _lifecycle_coordinator + if _user_sync_store is not None: + await _user_sync_store.close() if _nc is not None: await _nc.close() _nc = None diff --git a/app/operation/__init__.py b/app/operation/__init__.py index f078f2762..c2681a4a9 100644 --- a/app/operation/__init__.py +++ b/app/operation/__init__.py @@ -139,15 +139,34 @@ async def get_validated_host(self, db: AsyncSession, host_id: int) -> ProxyHost: await self.raise_error(message="Host not found", code=404) return db_host - async def get_validated_sub(self, db: AsyncSession, token: str, *, load_admin_role: bool = False) -> User: + async def get_validated_sub( + self, + db: AsyncSession, + token: str, + *, + load_admin: bool = True, + load_admin_role: bool = False, + load_next_plan: bool = True, + load_usage_logs: bool = True, + load_groups: bool = True, + load_lifetime_used_traffic: bool = False, + ) -> User: sub = await get_subscription_payload(token) db_user = None if sub: + load_kwargs = { + "load_admin": load_admin, + "load_admin_role": load_admin_role, + "load_next_plan": load_next_plan, + "load_usage_logs": load_usage_logs, + "load_groups": load_groups, + "load_lifetime_used_traffic": load_lifetime_used_traffic, + } if sub.get("user_id"): - db_user = await get_user_by_id(db, sub["user_id"], load_admin_role=load_admin_role) + db_user = await get_user_by_id(db, sub["user_id"], **load_kwargs) elif sub.get("username"): - db_user = await get_user(db, sub["username"], load_admin_role=load_admin_role) + db_user = await get_user(db, sub["username"], **load_kwargs) if ( not db_user diff --git a/app/operation/subscription.py b/app/operation/subscription.py index 4d345dce2..c5535ee7f 100644 --- a/app/operation/subscription.py +++ b/app/operation/subscription.py @@ -1,4 +1,5 @@ import re +from datetime import UTC, datetime as dt from json import dumps as json_dumps from typing import Any, ClassVar @@ -19,6 +20,7 @@ from app.models.subscription import SubscriptionUsageQuery from app.models.user import SubscriptionUserResponse, UsersResponseWithInbounds from app.settings import hwid_settings, subscription_settings +from app.subscription import sub_update_buffer as _sub_update_buffer # noqa: F401 # registers per-worker flush loop from app.subscription.share import ( apply_custom_format_variables, encode_title, @@ -87,6 +89,18 @@ class SubscriptionOperation(BaseOperation): _ENCODED_RULE_RESPONSE_HEADERS: ClassVar[set[str]] = {"announce", "profile-title"} + _SUB_CONFIG_LOAD: ClassVar[dict[str, bool]] = { + "load_next_plan": False, + "load_usage_logs": False, + "load_groups": False, + "load_lifetime_used_traffic": True, + } + _SUB_INFO_LOAD: ClassVar[dict[str, bool]] = { + "load_next_plan": True, + "load_usage_logs": False, + "load_groups": True, + "load_lifetime_used_traffic": True, + } @staticmethod async def validated_user(db_user: User) -> UsersResponseWithInbounds: @@ -200,6 +214,7 @@ def create_response_headers( } if extra_headers: headers.update(extra_headers) + headers["Cache-Control"] = "no-store" return headers @classmethod @@ -389,6 +404,11 @@ async def validate_and_register_hwid( existing_hwid = await get_user_hwid_by_value(db, user_id, x_hwid) if existing_hwid: + last_used = existing_hwid.last_used_at + if last_used is not None and last_used.tzinfo is None: + last_used = last_used.replace(tzinfo=UTC) + if last_used is not None and (dt.now(UTC) - last_used).total_seconds() < 300: + return await register_user_hwid(db, user_id, x_hwid, x_device_os, x_ver_os, x_device_model) return @@ -417,7 +437,7 @@ async def user_subscription( Provides a subscription link based on the user agent (Clash, V2Ray, etc.). """ sub_settings: SubSettings = await subscription_settings() - db_user = await self.get_validated_sub(db, token, load_admin_role=True) + db_user = await self.get_validated_sub(db, token, load_admin_role=True, **self._SUB_CONFIG_LOAD) role_hwid_settings = db_user.admin.role.hwid if db_user.admin and db_user.admin.role else None user = await self.validated_user(db_user) is_browser_request = "text/html" in accept_header @@ -540,7 +560,7 @@ async def user_subscription_with_client_type( if client_type == ConfigFormat.block or not getattr(sub_settings.manual_sub_request, client_type): await self.raise_error(message="Client not supported", code=406) - db_user = await self.get_validated_sub(db, token=token, load_admin_role=True) + db_user = await self.get_validated_sub(db, token=token, load_admin_role=True, **self._SUB_CONFIG_LOAD) user = await self.validated_user(db_user) await self.validate_and_register_hwid( @@ -612,7 +632,7 @@ def _build_raw_subscription_payload( async def user_subscription_raw(self, db: AsyncSession, token: str, request_url: str = ""): sub_settings: SubSettings = await subscription_settings() - db_user = await self.get_validated_sub(db, token, load_admin_role=True) + db_user = await self.get_validated_sub(db, token, load_admin_role=True, **self._SUB_CONFIG_LOAD) user = await self.validated_user(db_user) is_hwid_enabled = await self.is_user_hwid_enabled(db_user) @@ -682,7 +702,7 @@ async def user_subscription_info( ) -> tuple[SubscriptionUserResponse, dict]: """Retrieves detailed information about the user's subscription.""" sub_settings: SubSettings = await subscription_settings() - db_user = await self.get_validated_sub(db, token=token) + db_user = await self.get_validated_sub(db, token=token, **self._SUB_INFO_LOAD) user = await self.validated_user(db_user) response_headers = self.create_info_response_headers(user, sub_settings) @@ -699,7 +719,7 @@ async def user_subscription_apps(self, db: AsyncSession, token: str) -> list[App """ Get available applications for user's subscription. """ - db_user = await self.get_validated_sub(db, token=token, load_admin_role=True) + db_user = await self.get_validated_sub(db, token=token, load_admin_role=True, **self._SUB_CONFIG_LOAD) user = await self.validated_user(db_user) is_hwid_enabled = await self.is_user_hwid_enabled(db_user) sub_settings: SubSettings = await subscription_settings() @@ -738,7 +758,7 @@ async def user_subscription_headers( Retrieves only the headers for a subscription request, bypassing configuration generation. """ sub_settings: SubSettings = await subscription_settings() - db_user = await self.get_validated_sub(db, token, load_admin_role=True) + db_user = await self.get_validated_sub(db, token, load_admin_role=True, **self._SUB_CONFIG_LOAD) user = await self.validated_user(db_user) is_browser_request = "text/html" in accept_header is_subscription_page_request = is_browser_request and not sub_settings.disable_sub_template @@ -791,6 +811,13 @@ async def get_user_usage( """Fetches the usage statistics for the user within a specified date range.""" start, end = await self.validate_dates(query.start, query.end, True) - db_user = await self.get_validated_sub(db, token=token) + db_user = await self.get_validated_sub( + db, + token=token, + load_admin=False, + load_next_plan=False, + load_usage_logs=False, + load_groups=False, + ) return await get_user_usages(db, db_user.id, start, end, query.period) diff --git a/app/subscription/base.py b/app/subscription/base.py index 191667679..9512495db 100644 --- a/app/subscription/base.py +++ b/app/subscription/base.py @@ -56,6 +56,11 @@ def clean_dict(d: dict) -> dict: return clean_dict(data) +def dumps_compact(obj: Any) -> str: + """JSON for clients: no pretty-print whitespace.""" + return json.dumps(obj, separators=(",", ":"), ensure_ascii=False) + + class BaseSubscription: def __init__( self, diff --git a/app/subscription/clash.py b/app/subscription/clash.py index 6a8904b0b..84b9c83fb 100644 --- a/app/subscription/clash.py +++ b/app/subscription/clash.py @@ -1,7 +1,5 @@ from random import choice -from uuid import UUID -import yaml from pydantic import BaseModel from app.models.subscription import ( @@ -13,7 +11,6 @@ XHTTPTransportConfig, ) from app.templates import render_template_string -from app.utils.helpers import yml_uuid_representer from . import BaseSubscription @@ -60,16 +57,9 @@ def __init__( } def render(self): - yaml.add_representer(UUID, yml_uuid_representer) - return yaml.dump( - yaml.safe_load( - render_template_string( - self.clash_template_content, - {"conf": self.data, "proxy_remarks": self.proxy_remarks}, - ), - ), - sort_keys=False, - allow_unicode=True, + return render_template_string( + self.clash_template_content, + {"conf": self.data, "proxy_remarks": self.proxy_remarks}, ) def __str__(self) -> str: diff --git a/app/subscription/config_cache.py b/app/subscription/config_cache.py new file mode 100644 index 000000000..1c18c2c0c --- /dev/null +++ b/app/subscription/config_cache.py @@ -0,0 +1,63 @@ +"""Short in-process TTL cache for generated subscription payloads. + +Client apps poll /sub far more often than hosts or user settings change. +A few seconds of reuse avoids rebuilding JSON/YAML on every pull. Each +Uvicorn worker has its own cache (same model as the sub-update buffer). +""" + +from __future__ import annotations + +import time +from collections import OrderedDict +from typing import Any + +SUB_CONFIG_CACHE_TTL_S = 15 +SUB_CONFIG_CACHE_MAX = 4096 + +_cache: OrderedDict[tuple, tuple[float, str | bytes]] = OrderedDict() + + +def make_sub_config_key( + user: Any, + config_format: str, + as_base64: bool, + randomize_order: bool, +) -> tuple: + expire = getattr(user, "expire", None) + expire_key = expire.timestamp() if hasattr(expire, "timestamp") else expire + status = getattr(user, "status", None) + status_key = status.value if hasattr(status, "value") else status + inbounds = getattr(user, "inbounds", None) or () + return ( + getattr(user, "id", None), + config_format, + bool(as_base64), + bool(randomize_order), + status_key, + getattr(user, "data_limit", None), + expire_key, + tuple(inbounds), + ) + + +def get_sub_config(key: tuple) -> str | bytes | None: + item = _cache.get(key) + if item is None: + return None + expires_at, value = item + if expires_at <= time.monotonic(): + _cache.pop(key, None) + return None + _cache.move_to_end(key) + return value + + +def put_sub_config(key: tuple, value: str | bytes) -> None: + _cache[key] = (time.monotonic() + SUB_CONFIG_CACHE_TTL_S, value) + _cache.move_to_end(key) + while len(_cache) > SUB_CONFIG_CACHE_MAX: + _cache.popitem(last=False) + + +def clear_sub_config_cache() -> None: + _cache.clear() diff --git a/app/subscription/outline.py b/app/subscription/outline.py index aa3916a5a..d77c829f8 100644 --- a/app/subscription/outline.py +++ b/app/subscription/outline.py @@ -1,8 +1,6 @@ -import json - from app.models.subscription import SubscriptionInboundData -from .base import BaseSubscription +from .base import BaseSubscription, dumps_compact class OutlineConfiguration(BaseSubscription): @@ -13,7 +11,7 @@ def add_directly(self, data: dict): self.config.update(data) def render(self): - return json.dumps(self.config, indent=0) + return dumps_compact(self.config) def _build_shadowsocks(self, remark: str, address: str, inbound: SubscriptionInboundData, settings: dict) -> dict: """Build Shadowsocks outbound with 2022 support""" diff --git a/app/subscription/share.py b/app/subscription/share.py index 986bc98b4..beb37a8bf 100644 --- a/app/subscription/share.py +++ b/app/subscription/share.py @@ -14,6 +14,7 @@ from app.models.user import UsersResponseWithInbounds from app.settings import subscription_settings from app.subscription.client_templates import subscription_client_templates, subscription_xray_templates +from app.subscription.config_cache import get_sub_config, make_sub_config_key, put_sub_config from app.utils.system import readable_size from . import ( @@ -83,6 +84,11 @@ async def generate_subscription( as_base64: bool, randomize_order: bool = False, ) -> str | bytes: + cache_key = make_sub_config_key(user, config_format, as_base64, randomize_order) + cached = get_sub_config(cache_key) + if cached is not None: + return cached + client_templates = await subscription_client_templates() xray_template_overrides = await subscription_xray_templates() if config_format == "xray" else None conf = _build_subscription_config(config_format, client_templates) @@ -106,6 +112,7 @@ async def generate_subscription( if as_base64 and not isinstance(config, bytes): config = base64.b64encode(config.encode()).decode() + put_sub_config(cache_key, config) return config @@ -338,36 +345,28 @@ async def process_host( if inbound.use_sni_as_host and sni: req_host = sni - # Create a copy of the inbound data with selected random values - # Deep copy tls_config and transport_config to avoid mutating cached host data - inbound_copy = inbound.model_copy( - update={ - "tls_config": inbound.tls_config.model_copy(deep=True), - "transport_config": inbound.transport_config.model_copy(deep=True), - } - ) - - # Update TLS config with selected values - inbound_copy.tls_config.sni = sni - inbound_copy.tls_config.reality_short_id = reality_sid - - # Update transport config with selected host - inbound_copy.transport_config.host = req_host - inbound_copy.transport_config.path = path - if getattr(inbound_copy.transport_config, "request", None): - inbound_copy.transport_config.request = _format_dynamic_value( - inbound_copy.transport_config.request, + # Copy only the nested models we mutate so cached host objects stay intact. + transport_update: dict = {"host": req_host, "path": path} + if getattr(inbound.transport_config, "request", None): + transport_update["request"] = _format_dynamic_value( + inbound.transport_config.request, format_variables, ) - if getattr(inbound_copy.transport_config, "response", None): - inbound_copy.transport_config.response = _format_dynamic_value( - inbound_copy.transport_config.response, + if getattr(inbound.transport_config, "response", None): + transport_update["response"] = _format_dynamic_value( + inbound.transport_config.response, format_variables, ) - - # Update address and port with selected values - inbound_copy.address = address - inbound_copy.port = port + inbound_copy = inbound.model_copy( + update={ + "tls_config": inbound.tls_config.model_copy( + update={"sni": sni, "reality_short_id": reality_sid}, + ), + "transport_config": inbound.transport_config.model_copy(update=transport_update), + "address": address, + "port": port, + } + ) return inbound_copy, settings diff --git a/app/subscription/singbox.py b/app/subscription/singbox.py index a484965a1..4b91723d0 100644 --- a/app/subscription/singbox.py +++ b/app/subscription/singbox.py @@ -11,6 +11,7 @@ ) from . import BaseSubscription +from .base import dumps_compact class SingBoxConfiguration(BaseSubscription): @@ -58,7 +59,7 @@ def add_endpoint(self, endpoint_data): def render(self): self._finalize_config() - return json.dumps(self.config, indent=4) + return dumps_compact(self.config) def _finalize_config(self): urltest_types = ["vmess", "vless", "trojan", "shadowsocks", "hysteria2", "tuic", "http", "ssh"] diff --git a/app/subscription/sub_update_buffer.py b/app/subscription/sub_update_buffer.py new file mode 100644 index 000000000..28379bde8 --- /dev/null +++ b/app/subscription/sub_update_buffer.py @@ -0,0 +1,151 @@ +"""Buffer public subscription-update writes so GET /sub stays read-mostly. + +Each process keeps an in-memory queue and flushes on size, a short interval +(every worker — APScheduler is leader-only), shutdown, and before admin reads. +""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime as dt +from typing import Any + +from sqlalchemy import insert, select + +from app.db import GetDB +from app.db.models import User, UserSubscriptionUpdate +from app.lifecycle import on_shutdown, on_startup +from app.utils.logger import get_logger +from config import runtime_settings + +logger = get_logger("sub-update-buffer") + +FLUSH_INTERVAL_SECONDS = 2.0 +FLUSH_BATCH_SIZE = 100 +_MAX_BUFFER = 50_000 + +_USER_AGENT_MAX_LEN = UserSubscriptionUpdate.__table__.columns.user_agent.type.length or 512 +_IP_MAX_LEN = UserSubscriptionUpdate.__table__.columns.ip.type.length or 64 +_HWID_MAX_LEN = UserSubscriptionUpdate.__table__.columns.hwid.type.length or 256 + +_pending: list[dict[str, Any]] = [] +_lock = asyncio.Lock() +_drain_lock = asyncio.Lock() +_flush_task: asyncio.Task | None = None + + +def _sanitize_record(user_id: int, user_agent: str, ip: str | None, hwid: str | None) -> dict[str, Any]: + sanitized_ip = (ip or "")[:_IP_MAX_LEN] or None + sanitized_hwid = (hwid or "")[:_HWID_MAX_LEN] or None + return { + "user_id": user_id, + "user_agent": (user_agent or "")[:_USER_AGENT_MAX_LEN], + "ip": sanitized_ip, + "hwid": sanitized_hwid, + "created_at": dt.now(UTC), + } + + +def pending_count() -> int: + return len(_pending) + + +async def reset_user_sub_update_buffer() -> None: + """Drop queued rows without writing. Tests only.""" + async with _lock: + _pending.clear() + + +async def queue_user_sub_update(user_id: int, user_agent: str, ip: str | None = None, hwid: str | None = None) -> None: + """Enqueue a subscription-update row; may kick a background flush.""" + record = _sanitize_record(user_id, user_agent, ip, hwid) + should_flush = False + dropped = 0 + async with _lock: + prev = len(_pending) + _pending.append(record) + overflow = len(_pending) - _MAX_BUFFER + if overflow > 0: + del _pending[:overflow] + dropped = overflow + should_flush = prev < FLUSH_BATCH_SIZE <= len(_pending) + if dropped: + logger.warning("Dropped %s buffered subscription updates; buffer full", dropped) + if should_flush: + asyncio.create_task(flush_user_sub_updates(), name="sub_update_flush") + + +async def flush_user_sub_updates() -> int: + """Persist queued rows. Concurrent callers wait, then drain anything left.""" + async with _drain_lock: + written = 0 + while True: + async with _lock: + if not _pending: + return written + batch = _pending[:FLUSH_BATCH_SIZE] + del _pending[:FLUSH_BATCH_SIZE] + try: + async with GetDB() as db: + # Users can be deleted after their subscription request was queued. + # Lock surviving parents until commit so concurrent deletes cannot + # invalidate the foreign keys between this read and the insert. + user_ids = sorted({record["user_id"] for record in batch}) + existing_user_ids = set( + await db.scalars( + select(User.id) + .where(User.id.in_(user_ids)) + .order_by(User.id) + .with_for_update(read=True, key_share=True) + ) + ) + live_records = [record for record in batch if record["user_id"] in existing_user_ids] + if live_records: + await db.execute(insert(UserSubscriptionUpdate.__table__), live_records) + await db.commit() + written += len(live_records) + except Exception: + async with _lock: + _pending[0:0] = batch + overflow = len(_pending) - _MAX_BUFFER + if overflow > 0: + del _pending[:overflow] + logger.exception("Failed to flush %s buffered subscription updates", len(batch)) + raise + + +async def _flush_loop() -> None: + while True: + await asyncio.sleep(FLUSH_INTERVAL_SECONDS) + try: + await flush_user_sub_updates() + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Periodic subscription-update flush failed") + + +@on_startup +async def start_sub_update_flusher() -> None: + global _flush_task + if not runtime_settings.role.runs_panel: + return + if _flush_task is not None and not _flush_task.done(): + return + _flush_task = asyncio.create_task(_flush_loop(), name="sub_update_flush_loop") + + +@on_shutdown +async def stop_sub_update_flusher() -> None: + global _flush_task + if _flush_task is not None: + _flush_task.cancel() + try: + await _flush_task + except asyncio.CancelledError: + pass + _flush_task = None + try: + await flush_user_sub_updates() + except Exception: + logger.exception("Final subscription-update flush failed") diff --git a/app/subscription/xray.py b/app/subscription/xray.py index aa20452f6..3f6c3d740 100644 --- a/app/subscription/xray.py +++ b/app/subscription/xray.py @@ -15,6 +15,7 @@ ) from . import BaseSubscription +from .base import dumps_compact class XrayConfiguration(BaseSubscription): @@ -63,7 +64,7 @@ def add_config(self, remarks, outbounds, template_content: str | None = None): self.config.append(json_template) def render(self): - return json.dumps(self.config, indent=4) + return dumps_compact(self.config) def add( self, diff --git a/app/templates/filters.py b/app/templates/filters.py index 70713657d..0be1f0fa5 100644 --- a/app/templates/filters.py +++ b/app/templates/filters.py @@ -1,8 +1,10 @@ import os from datetime import UTC, datetime +from uuid import UUID import yaml +from app.utils.helpers import yml_uuid_representer from app.utils.system import readable_size @@ -10,7 +12,8 @@ def to_yaml(obj): if not obj: return "" - return yaml.dump(obj, allow_unicode=True, indent=2) + yaml.add_representer(UUID, yml_uuid_representer) + return yaml.dump(obj, allow_unicode=True, indent=2, sort_keys=False) def exclude_keys(obj, *target_keys): diff --git a/tests/api/__init__.py b/tests/api/__init__.py index 71100544c..f5daf0a97 100644 --- a/tests/api/__init__.py +++ b/tests/api/__init__.py @@ -1,5 +1,6 @@ import asyncio import json +import logging from alembic.command import upgrade from alembic.config import Config @@ -78,9 +79,20 @@ def run_migrations_sync(): alembic_cfg = Config("alembic.ini") alembic_cfg.set_main_option("sqlalchemy.url", sync_db_url) - with create_engine(sync_db_url).begin() as connection: - alembic_cfg.attributes["connection"] = connection - upgrade(alembic_cfg, "head") + # Alembic's fileConfig disables pre-existing application loggers. Keep + # migration setup from making later log assertions depend on test order. + logger_states = { + logger: logger.disabled + for logger in logging.root.manager.loggerDict.values() + if isinstance(logger, logging.Logger) + } + try: + with create_engine(sync_db_url).begin() as connection: + alembic_cfg.attributes["connection"] = connection + upgrade(alembic_cfg, "head") + finally: + for logger, disabled in logger_states.items(): + logger.disabled = disabled print("[migrations] Migrations completed successfully") return True # Migrations ran successfully except Exception as e: diff --git a/tests/api/conftest.py b/tests/api/conftest.py index 16a406b90..56febc18c 100644 --- a/tests/api/conftest.py +++ b/tests/api/conftest.py @@ -25,6 +25,7 @@ def mock_db_session(monkeypatch: pytest.MonkeyPatch): db_session = MagicMock(spec=TestSession) monkeypatch.setattr("app.settings.GetDB", db_session) monkeypatch.setattr("app.subscription.client_templates.GetDB", GetTestDB) + monkeypatch.setattr("app.subscription.sub_update_buffer.GetDB", GetTestDB) return db_session diff --git a/tests/api/test_subscription_update_snapshot.py b/tests/api/test_subscription_update_snapshot.py new file mode 100644 index 000000000..4d88aa483 --- /dev/null +++ b/tests/api/test_subscription_update_snapshot.py @@ -0,0 +1,64 @@ +from datetime import UTC, datetime, timedelta + +import pytest +from sqlalchemy import delete, select + +from app.db.crud.user import ( + get_users_sub_update_list, + get_users_subscription_agent_counts, + get_users_subscription_agent_stats, +) +from app.db.models import User, UserSubscriptionUpdate +from app.models.stats import Period +from app.subscription import sub_update_buffer +from tests.api import TestSession, engine +from tests.api.helpers import unique_name + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reader", ["list", "counts", "stats"]) +@pytest.mark.parametrize("already_flushed", [False, True]) +async def test_subscription_reads_see_updates_after_request_snapshot(reader, already_flushed): + async with TestSession() as setup: + user = User(username=unique_name("sub_snapshot")) + setup.add(user) + await setup.commit() + user_id = user.id + + try: + async with TestSession() as request: + if engine.dialect.name in {"mysql", "mariadb", "postgresql"}: + await request.connection(execution_options={"isolation_level": "REPEATABLE READ"}) + # Authorization/user lookup establishes the request's read snapshot. + db_user = await request.scalar(select(User).where(User.id == user_id)) + db_user.note = "uncommitted caller change" + await sub_update_buffer.queue_user_sub_update(user_id, "snapshot-client") + if already_flushed: + await sub_update_buffer.flush_user_sub_updates() + + start = datetime.now(UTC).replace(minute=0, second=0, microsecond=0) - timedelta(hours=1) + end = start + timedelta(hours=3) + if reader == "list": + rows, count = await get_users_sub_update_list(request, user_id) + assert count == 1 + assert [row.user_agent for row in rows] == ["snapshot-client"] + elif reader == "counts": + counts = await get_users_subscription_agent_counts(request, user_id=user_id, start=start, end=end) + assert counts == [("snapshot-client", 1)] + else: + stats = await get_users_subscription_agent_stats(request, start, end, Period.hour, user_id=user_id) + assert [(row["agent"], row["count"]) for row in stats] == [("snapshot-client", 1)] + + assert request.in_transaction() + assert db_user in request.dirty + assert db_user.note == "uncommitted caller change" + + async with TestSession() as check: + assert await check.scalar(select(User.note).where(User.id == user_id)) is None + finally: + await sub_update_buffer.flush_user_sub_updates() + async with TestSession() as cleanup: + # The API SQLite test engine does not enable foreign-key cascades. + await cleanup.execute(delete(UserSubscriptionUpdate).where(UserSubscriptionUpdate.user_id == user_id)) + await cleanup.execute(delete(User).where(User.id == user_id)) + await cleanup.commit() diff --git a/tests/nats_sync_process_worker.py b/tests/nats_sync_process_worker.py new file mode 100644 index 000000000..4922018cf --- /dev/null +++ b/tests/nats_sync_process_worker.py @@ -0,0 +1,40 @@ +"""Independent Python worker used by the real NATS integration tests.""" + +import asyncio +import json +import sys +from pathlib import Path + +import nats + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from app.node.nats_memory import NatsUserSyncStore + + +async def main(): + url, bucket, worker_id, node_count = sys.argv[1:] + nc = await nats.connect(url) + store = NatsUserSyncStore(await nc.jetstream().key_value(bucket)) + emails = [] + try: + async with asyncio.timeout(60): + # All inputs are seeded before process startup. An empty complete + # pass means remaining work belongs to one of the other processes. + while True: + progress = False + for node in range(int(node_count)): + claimed = await store.claim_users(str(node), worker_id, 50, 120) + emails.extend(item.user.email for item in claimed) + await store.ack_users(str(node), [item.token for item in claimed]) + progress |= bool(claimed) + if not progress: + break + print(json.dumps(emails), flush=True) + finally: + await store.close() + await nc.close() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/test_nats_kv_index.py b/tests/test_nats_kv_index.py new file mode 100644 index 000000000..f71429767 --- /dev/null +++ b/tests/test_nats_kv_index.py @@ -0,0 +1,194 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from app.nats.kv_index import KvKeyIndex + + +class Watcher: + def __init__(self, entries=()): + self.queue = asyncio.Queue() + for entry in entries: + self.queue.put_nowait(entry) + self.stopped = False + + def __aiter__(self): + return self + + async def __anext__(self): + entry = await self.queue.get() + if isinstance(entry, Exception): + raise entry + return entry + + async def stop(self): + self.stopped = True + + +def entry(key, revision=1, operation=None): + return SimpleNamespace(key=key, revision=revision, operation=operation) + + +class WatchedKv: + def __init__(self, *watchers): + self.watchers = iter(watchers) + self.calls = 0 + + async def watch(self, subject, **kwargs): + assert subject == ">" + assert kwargs["meta_only"] is True + assert not kwargs.get("ignore_deletes") + self.calls += 1 + return next(self.watchers) + + +async def eventually(predicate): + async with asyncio.timeout(3): + while not predicate(): + await asyncio.sleep(0.001) + + +async def test_many_nodes_and_concurrent_empty_polls_share_one_watcher(): + kv = WatchedKv(Watcher([None])) + index = KvKeyIndex(kv) + try: + assert await asyncio.gather(*(index.keys(f"p.{node}.") for node in range(1000))) == [[]] * 1000 + assert kv.calls == 1 + assert index._keys == {} + finally: + await index.close() + + +async def test_initial_snapshot_is_complete_before_exposing_keys(): + watcher = Watcher([entry("p.1.a"), entry("p.2.b")]) + index = KvKeyIndex(WatchedKv(watcher)) + task = asyncio.create_task(index.keys("p.1.")) + try: + await eventually(lambda: "p.2." in index._keys) + assert not task.done() + watcher.queue.put_nowait(None) + assert await task == ["p.1.a"] + assert await index.keys("p.2.") == ["p.2.b"] + finally: + await index.close() + + +async def test_slow_snapshot_can_finish_while_replay_keeps_progressing(monkeypatch): + monkeypatch.setattr(KvKeyIndex, "SNAPSHOT_STALL_TIMEOUT", 0.1) + watcher = Watcher() + index = KvKeyIndex(WatchedKv(watcher)) + pending = asyncio.create_task(index.keys("p.1.")) + try: + for revision in range(1, 13): + watcher.queue.put_nowait(entry("p.1.a", revision)) + await asyncio.sleep(0.02) + assert not pending.done() + watcher.queue.put_nowait(None) + assert await pending == ["p.1.a"] + finally: + await index.close() + + +async def test_stalled_snapshot_still_times_out(monkeypatch): + monkeypatch.setattr(KvKeyIndex, "SNAPSHOT_STALL_TIMEOUT", 0.05) + index = KvKeyIndex(WatchedKv(Watcher())) + try: + with pytest.raises(TimeoutError): + await index.keys("p.1.") + finally: + await index.close() + + +async def test_changes_remove_deleted_keys_and_empty_node_indexes(): + watcher = Watcher([entry("p.1.a"), None]) + index = KvKeyIndex(WatchedKv(watcher)) + try: + assert await index.keys("p.1.") == ["p.1.a"] + for number in range(1000): + watcher.queue.put_nowait(entry(f"c.1.{number}", number + 2)) + watcher.queue.put_nowait(entry(f"c.1.{number}", number + 3, "DEL")) + watcher.queue.put_nowait(entry("p.1.a", 3000, "PURGE")) + await eventually(lambda: watcher.queue.empty()) + assert index._keys == {} + finally: + await index.close() + assert watcher.stopped + + +async def test_watch_failure_rebuilds_snapshot_without_stale_keys(): + first = Watcher([entry("p.1.old"), None]) + second = Watcher([entry("p.1.new", 2), None]) + kv = WatchedKv(first, second) + index = KvKeyIndex(kv) + try: + assert await index.keys("p.1.") == ["p.1.old"] + first.queue.put_nowait(RuntimeError("connection lost")) + await eventually(lambda: first.stopped) + assert not index._ready.is_set() + assert await index.keys("p.1.") == ["p.1.new"] + assert kv.calls == 2 + finally: + await index.close() + + +async def test_cancelled_caller_does_not_cancel_shared_initialization(): + watcher = Watcher() + index = KvKeyIndex(WatchedKv(watcher)) + cancelled = asyncio.create_task(index.keys("p.1.")) + remaining = asyncio.create_task(index.keys("p.2.")) + try: + await asyncio.sleep(0) + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + watcher.queue.put_nowait(entry("p.2.a")) + watcher.queue.put_nowait(None) + assert await remaining == ["p.2.a"] + finally: + await index.close() + + +async def test_close_wakes_waiters_and_rejects_future_reads(): + index = KvKeyIndex(WatchedKv(Watcher())) + pending = asyncio.create_task(index.keys("p.1.")) + await asyncio.sleep(0) + await index.close() + with pytest.raises(RuntimeError, match="closed"): + await pending + with pytest.raises(RuntimeError, match="closed"): + await index.keys("p.1.") + + +async def test_local_writes_are_visible_before_delayed_watch_events(): + watcher = Watcher([None]) + index = KvKeyIndex(WatchedKv(watcher)) + try: + assert await index.keys("p.1.") == [] + index.observe_put("p.1.a", 3) + index.observe_put("p.1.a", 1) # a delayed acknowledgement cannot regress it + watcher.queue.put_nowait(entry("p.1.a", 2, "DEL")) + await eventually(lambda: watcher.queue.empty()) + assert await index.entries("p.1.") == {"p.1.a": 3} + watcher.queue.put_nowait(entry("p.1.a", 3)) + await eventually(lambda: watcher.queue.empty()) + assert index._local_puts == {} + watcher.queue.put_nowait(entry("p.1.a", 4, "DEL")) + await eventually(lambda: watcher.queue.empty()) + index.observe_put("p.1.a", 3) # delete already observed before publish reply + assert await index.keys("p.1.") == [] + finally: + await index.close() + + +async def test_missing_candidate_eviction_does_not_remove_a_newer_update(): + index = KvKeyIndex(WatchedKv(Watcher([entry("p.1.a", 1), None]))) + try: + assert await index.entries("p.1.") == {"p.1.a": 1} + index.observe_put("p.1.a", 2) + index.discard("p.1.a", 1) + assert await index.entries("p.1.") == {"p.1.a": 2} + index.discard("p.1.a", 2) + assert index._keys == {} + finally: + await index.close() diff --git a/tests/test_nats_node_memory.py b/tests/test_nats_node_memory.py index bbf0d9fa7..4f22af923 100644 --- a/tests/test_nats_node_memory.py +++ b/tests/test_nats_node_memory.py @@ -72,6 +72,31 @@ async def test_user_sync_enqueue_shards_per_email_key(): assert kv._data == {} +@pytest.mark.asyncio +async def test_bulk_sync_bounds_concurrent_writes_across_nodes(): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + create = kv.create + active = peak = 0 + + async def slow_create(key, value): + nonlocal active, peak + active += 1 + peak = max(peak, active) + try: + await asyncio.sleep(0.001) + return await create(key, value) + finally: + active -= 1 + + kv.create = slow_create + users = [_user(f"user{i}") for i in range(200)] + await asyncio.gather(*(store.enqueue_users(str(node), users) for node in range(4))) + assert len(kv._data) == 800 + assert 1 < peak <= 32 + assert active == 0 + + @pytest.mark.asyncio async def test_claim_cleans_up_claimed_key_when_pending_delete_fails(): kv = MemoryCasKv() diff --git a/tests/test_nats_sync_integration.py b/tests/test_nats_sync_integration.py new file mode 100644 index 000000000..b1d5c96b9 --- /dev/null +++ b/tests/test_nats_sync_integration.py @@ -0,0 +1,503 @@ +"""NATS user synchronization integration tests. Set NATS_SERVER_BINARY or install nats-server.""" + +import asyncio +import contextlib +import json +import os +import shutil +import socket +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace +from uuid import uuid4 + +import nats +import pytest +import pytest_asyncio +from nats.js.errors import KeyNotFoundError +from PasarGuardNodeBridge.common.service_pb2 import User + +from app.nats.kv_cleanup import compact_deleted_keys +from app.nats.kv_watch import KvWatcher, watch_kv +from app.node.nats_memory import NatsUserSyncStore + + +@pytest_asyncio.fixture(loop_scope="function") +async def jetstream(tmp_path): + binary = os.environ.get("NATS_SERVER_BINARY") or shutil.which("nats-server") + if not binary: + pytest.skip("Set NATS_SERVER_BINARY to run JetStream integration tests") + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + process = await asyncio.to_thread( + subprocess.Popen, + [binary, "-js", "-a", "127.0.0.1", "-p", str(port), "-sd", str(tmp_path / "nats")], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0, + ) + nc = None + try: + async with asyncio.timeout(10): + while True: + try: + _reader, writer = await asyncio.open_connection("127.0.0.1", port) + writer.close() + await writer.wait_closed() + break + except OSError: + if process.poll() is not None: + raise RuntimeError("Test NATS server exited during startup") + await asyncio.sleep(0.02) + url = f"nats://127.0.0.1:{port}" + nc = await nats.connect(url) + js = nc.jetstream() + bucket = "sync_" + uuid4().hex + kv = await js.create_key_value(bucket=bucket) + + async def restart(): + nonlocal process + process.terminate() + await asyncio.to_thread(process.wait, timeout=10) + process = await asyncio.to_thread( + subprocess.Popen, + process.args, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0, + ) + async with asyncio.timeout(10): + while not nc.is_connected or nc.stats["reconnects"] == 0: + await asyncio.sleep(0.02) + + yield SimpleNamespace(nc=nc, js=js, kv=kv, bucket=bucket, process=process, url=url, restart=restart) + finally: + if nc is not None: + await nc.close() + process.terminate() + await asyncio.to_thread(process.wait, timeout=10) + + +async def wait_for_keys(store, prefix, count, timeout=10): + async with asyncio.timeout(timeout): + while len(await store._key_index.keys(prefix)) != count: + await asyncio.sleep(0.005) + + +@pytest.mark.parametrize("snapshot_only", [False, True]) +async def test_snapshot_pauses_delivery_until_reader_requests_next_batch(jetstream, monkeypatch, snapshot_only): + monkeypatch.setattr(KvWatcher, "BATCH_SIZE", 8) + expected = {} + for number in range(100): + key = f"p.1.{number}" + await jetstream.kv.put(key, b"value") + expected[key] = "DEL" if number % 2 else None + if number % 2: + await jetstream.kv.delete(key) + watcher = await watch_kv(jetstream.kv, ">", snapshot_only=snapshot_only) + try: + first = await anext(watcher) + delivered = (await watcher._snapshot.consumer_info()).delivered.consumer_seq + await asyncio.sleep(0.1) + assert (await watcher._snapshot.consumer_info()).delivered.consumer_seq == delivered == 8 + assert watcher._snapshot.pending_msgs == 0 + actual = {first.key: first.operation} + async for entry in watcher: + if entry is None: + break + assert entry.key not in actual + actual[entry.key] = entry.operation + assert actual == expected + info = await jetstream.js.stream_info(f"KV_{jetstream.bucket}") + assert info.state.consumer_count == (0 if snapshot_only else 1) + finally: + await watcher.stop() + assert (await jetstream.js.stream_info(f"KV_{jetstream.bucket}")).state.consumer_count == 0 + + +async def test_snapshot_to_live_handoff_preserves_concurrent_changes(jetstream, monkeypatch): + await jetstream.kv.put("p.1.existing", b"old") + subscribe = jetstream.js.subscribe + revisions = {} + + async def write_before_live_subscription(*args, **kwargs): + revisions["p.1.existing"] = await jetstream.kv.put("p.1.existing", b"new") + revisions["p.1.added"] = await jetstream.kv.put("p.1.added", b"new") + await jetstream.kv.delete("p.1.existing") + return await subscribe(*args, **kwargs) + + monkeypatch.setattr(jetstream.js, "subscribe", write_before_live_subscription) + watcher = await watch_kv(jetstream.kv, ">") + try: + assert (await anext(watcher)).key == "p.1.existing" + assert await anext(watcher) is None + async with asyncio.timeout(5): + entries = [await anext(watcher) for _ in range(2)] + assert [(entry.key, entry.operation) for entry in entries] == [ + ("p.1.added", None), + ("p.1.existing", "DEL"), + ] + # A history-one bucket retains the final delete instead of the + # intermediate put. Both keys' latest state must reach the index. + assert entries[0].revision == revisions["p.1.added"] + assert entries[1].revision > revisions["p.1.existing"] + finally: + await watcher.stop() + + +async def test_interrupted_snapshot_is_rejected_and_consumer_is_removed(jetstream, monkeypatch): + monkeypatch.setattr(KvWatcher, "BATCH_SIZE", 1) + for number in range(3): + await jetstream.kv.put(f"p.1.{number}", b"value") + nc = await nats.connect(jetstream.url, reconnect_time_wait=0.05) + watcher = await watch_kv(await nc.jetstream().key_value(jetstream.bucket), ">") + try: + await anext(watcher) + nc._transport.close() + async with asyncio.timeout(5): + while nc.stats["reconnects"] == 0: + await asyncio.sleep(0.02) + with pytest.raises(RuntimeError, match="connection changed during KV snapshot"): + await anext(watcher) + assert (await jetstream.js.stream_info(f"KV_{jetstream.bucket}")).state.consumer_count == 0 + finally: + await watcher.stop() + await nc.close() + + +async def test_stopping_partial_snapshot_removes_consumer(jetstream): + for number in range(3): + await jetstream.kv.put(f"p.1.{number}", b"value") + watcher = await watch_kv(jetstream.kv, ">") + await anext(watcher) + await watcher.stop() + assert (await jetstream.js.stream_info(f"KV_{jetstream.bucket}")).state.consumer_count == 0 + + +async def test_idle_polls_do_not_create_consumers_or_send_requests(jetstream): + stores = [NatsUserSyncStore(jetstream.kv) for _ in range(4)] + try: + # Seed deletion markers for completed claims. + for number in range(300): + await jetstream.kv.put(f"c.{number % 8}.{number}", b"{}") + await jetstream.kv.delete(f"c.{number % 8}.{number}") + for store in stores: + assert await store.claim_users("0", "warmup", 10, 30) == [] + before = await jetstream.js.stream_info(f"KV_{jetstream.bucket}") + sent = jetstream.nc.stats["out_msgs"] + for _ in range(20): + results = await asyncio.gather( + *( + store.claim_users(str(node), str(worker), 10, 30) + for worker, store in enumerate(stores) + for node in range(8) + ) + ) + assert all(result == [] for result in results) + assert jetstream.nc.stats["out_msgs"] == sent + after = await jetstream.js.stream_info(f"KV_{jetstream.bucket}") + assert before.state.consumer_count == after.state.consumer_count == 4 + assert all(store._key_index._keys == {} for store in stores) + finally: + await asyncio.gather(*(store.close() for store in stores)) + + +async def test_repeated_clear_does_not_replay_old_deletion_markers(jetstream): + for number in range(1000): + await jetstream.kv.delete(f"c.1.{number}") + store = NatsUserSyncStore(jetstream.kv) + try: + assert await store._key_index.keys("c.1.") == [] + received = jetstream.nc.stats["in_msgs"] + for _ in range(20): + await store.clear("1") + assert jetstream.nc.stats["in_msgs"] - received < 400 + assert (await jetstream.js.stream_info(f"KV_{jetstream.bucket}")).state.consumer_count == 1 + finally: + await store.close() + + +async def test_clear_catches_unobserved_writes_and_preserves_other_nodes(jetstream, monkeypatch): + gate = asyncio.Event() + gate.set() + real_watch = watch_kv + + class DelayedWatcher: + def __init__(self, watcher): + self.watcher = watcher + + def __aiter__(self): + return self + + async def __anext__(self): + entry = await anext(self.watcher) + if entry is not None: + await gate.wait() + return entry + + async def stop(self): + await self.watcher.stop() + + async def delayed_watch(*args, **kwargs): + return DelayedWatcher(await real_watch(*args, **kwargs)) + + monkeypatch.setattr("app.nats.kv_index.watch_kv", delayed_watch) + kv = jetstream.kv + for key in ("p.1.cached", "c.1.cached"): + await kv.put(key, b"{}") + store = NatsUserSyncStore(kv) + try: + assert await store._key_index.keys("p.1.") == ["p.1.cached"] + gate.clear() + for key in ("p.1.new", "c.1.new", "p.2.keep"): + await kv.put(key, b"{}") + await kv.delete("p.1.cached") + await kv.put("p.1.cached", b"{}") + await store.enqueue_users("1", [User(email="local")]) + await store.clear("1") + for key in ("p.1.cached", "c.1.cached", "p.1.new", "c.1.new", store._pending_key("1", "local")): + with pytest.raises(KeyNotFoundError): + await kv.get(key) + assert (await kv.get("p.2.keep")).value == b"{}" + assert (await jetstream.js.stream_info(f"KV_{jetstream.bucket}")).state.consumer_count == 1 + finally: + await store.close() + + +async def test_concurrent_workers_claim_each_update_once(jetstream): + stores = [NatsUserSyncStore(jetstream.kv) for _ in range(4)] + users = [User(email=f"user{i}", inbounds=["in"]) for i in range(80)] + try: + await stores[0].enqueue_users("1", users) + await asyncio.gather(*(wait_for_keys(store, "p.1.", 80) for store in stores)) + batches = await asyncio.gather( + *(store.claim_users("1", str(worker), 80, 30) for worker, store in enumerate(stores)) + ) + emails = [item.user.email for batch in batches for item in batch] + assert len(emails) == len(set(emails)) == 80 + await asyncio.gather( + *(store.ack_users("1", [item.token for item in batch]) for store, batch in zip(stores, batches)) + ) + for store in stores: + await wait_for_keys(store, "p.1.", 0) + await wait_for_keys(store, "c.1.", 0) + assert await store.claim_users("1", "idle", 80, 30) == [] + finally: + await asyncio.gather(*(store.close() for store in stores)) + + +async def test_live_enqueue_and_expired_claim_recovery(jetstream): + first = NatsUserSyncStore(jetstream.kv) + second = NatsUserSyncStore(jetstream.kv) + try: + assert await second.claim_users("1", "other", 10, 30) == [] + await first.enqueue_users("1", [User(email="recover", inbounds=["in"])]) + await wait_for_keys(second, "p.1.", 1) + claimed = await first.claim_users("1", "crashed", 10, 0) + assert len(claimed) == 1 + await wait_for_keys(second, "p.1.", 0) + await wait_for_keys(second, "c.1.", 1) + recovered = await second.claim_users("1", "replacement", 10, 30) + assert [item.user.email for item in recovered] == ["recover"] + finally: + await first.close() + await second.close() + + +async def test_enqueue_is_claimable_even_when_watch_notifications_are_delayed(jetstream, monkeypatch): + gate = asyncio.Event() + real_watch = watch_kv + + class DelayedWatcher: + def __init__(self, watcher): + self.watcher = watcher + + def __aiter__(self): + return self + + async def __anext__(self): + update = await self.watcher.__anext__() + if update is not None: + await gate.wait() + return update + + async def stop(self): + await self.watcher.stop() + + async def delayed_watch(*args, **kwargs): + return DelayedWatcher(await real_watch(*args, **kwargs)) + + monkeypatch.setattr("app.nats.kv_index.watch_kv", delayed_watch) + store = NatsUserSyncStore(jetstream.kv) + try: + assert await store.claim_users("1", "local", 10, 30) == [] + await store.enqueue_users("1", [User(email="immediate")]) + claimed = await store.claim_users("1", "local", 10, 0) + assert [item.user.email for item in claimed] == ["immediate"] + recovered = await store.claim_users("1", "local", 10, 30) + assert [item.user.email for item in recovered] == ["immediate"] + assert recovered[0].token != claimed[0].token + claimed = recovered + await store.requeue_users("1", claimed) + assert [item.user.email for item in await store.claim_users("1", "local", 10, 30)] == ["immediate"] + finally: + await store.close() + + +async def test_compaction_keeps_pending_and_concurrently_recreated_values(jetstream): + kv = jetstream.kv + await kv.put("p.1.live", b"pending") + await kv.put("p.1.race", b"old") + await kv.delete("p.1.race") + await kv.put("c.1.completed", b"old claim") + await kv.delete("c.1.completed") + assert await compact_deleted_keys(jetstream.js, jetstream.bucket) == 0 # grace period + original_purge = jetstream.js.purge_stream + + async def recreate_before_purge(stream, **kwargs): + if kwargs["subject"].endswith("p.1.race"): + await kv.put("p.1.race", b"new update") + return await original_purge(stream, **kwargs) + + jetstream.js.purge_stream = recreate_before_purge + assert await compact_deleted_keys(jetstream.js, jetstream.bucket, older_than=0) == 2 + assert (await kv.get("p.1.live")).value == b"pending" + assert (await kv.get("p.1.race")).value == b"new update" + with pytest.raises(KeyNotFoundError): + await kv.get("c.1.completed") + assert (await jetstream.js.stream_info(f"KV_{jetstream.bucket}")).state.messages == 2 + + +async def test_watcher_resumes_after_network_disconnect(jetstream): + nc = await nats.connect(jetstream.url, reconnect_time_wait=0.05) + store = NatsUserSyncStore(await nc.jetstream().key_value(jetstream.bucket)) + try: + assert await store.claim_users("1", "reader", 10, 30) == [] + # Drop only this client's TCP transport; the writer remains connected. + nc._transport.close() + await jetstream.kv.put("p.1.new", b'{"email":"new","user":"CgNuZXc="}') + async with asyncio.timeout(10): + while nc.stats["reconnects"] == 0: + await asyncio.sleep(0.02) + await wait_for_keys(store, "p.1.", 1) + assert [item.user.email for item in await store.claim_users("1", "reader", 10, 30)] == ["new"] + finally: + await store.close() + with contextlib.suppress(Exception): + await nc.close() + + +async def test_four_python_processes_deliver_every_update_once(jetstream): + writer = NatsUserSyncStore(jetstream.kv) + expected = {f"node{node}-user{user}" for node in range(4) for user in range(250)} + for node in range(4): + await writer.enqueue_users(str(node), [User(email=f"node{node}-user{user}") for user in range(250)]) + processes = [] + try: + for worker in range(4): + processes.append( + await asyncio.create_subprocess_exec( + sys.executable, + str(Path(__file__).with_name("nats_sync_process_worker.py")), + jetstream.url, + jetstream.bucket, + str(worker), + "4", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0, + ) + ) + async with asyncio.timeout(75): + results = await asyncio.gather(*(process.communicate() for process in processes)) + emails = [] + for process, (stdout, stderr) in zip(processes, results): + assert process.returncode == 0, stderr.decode(errors="replace") + emails.extend(json.loads(stdout)) + assert len(emails) == len(expected) + assert set(emails) == expected + assert await writer.claim_users("0", "verify", 50, 120) == [] + await compact_deleted_keys(jetstream.js, jetstream.bucket, older_than=0) + assert (await jetstream.js.stream_info(f"KV_{jetstream.bucket}")).state.messages == 0 + finally: + for process in processes: + if process.returncode is None: + process.kill() + await process.wait() + await writer.close() + + +async def test_server_restart_preserves_pending_work_and_watch_updates(jetstream): + store = NatsUserSyncStore(jetstream.kv) + try: + await store.enqueue_users("1", [User(email="before-restart")]) + await wait_for_keys(store, "p.1.", 1) + await jetstream.restart() + # A separate store writes after reconnect so the reader must receive a + # watch notification, rather than using its own immediate-write hint. + writer = NatsUserSyncStore(jetstream.kv) + try: + await writer.enqueue_users("1", [User(email="after-restart")]) + finally: + await writer.close() + # nats-py checks heartbeat activity every 10 seconds and the first + # check clears the initial active flag. Allow two activity checks. + await wait_for_keys(store, "p.1.", 2, timeout=30) + claimed = await store.claim_users("1", "restarted", 10, 30) + assert {item.user.email for item in claimed} == {"before-restart", "after-restart"} + await store.ack_users("1", [item.token for item in claimed]) + await wait_for_keys(store, "c.1.", 0) + finally: + await store.close() + + +async def test_repeated_sync_compaction_and_idle_keep_resources_bounded(jetstream): + store = NatsUserSyncStore(jetstream.kv) + stream = f"KV_{jetstream.bucket}" + try: + for cycle in range(20): + users = [User(email=f"cycle{cycle}-user{number}") for number in range(50)] + await store.enqueue_users("1", users) + claimed = await store.claim_users("1", "worker", 100, 30) + assert {item.user.email for item in claimed} == {user.email for user in users} + await store.ack_users("1", [item.token for item in claimed]) + await wait_for_keys(store, "p.1.", 0) + await wait_for_keys(store, "c.1.", 0) + assert store._key_index._keys == {} + assert store._key_index._local_puts == {} + await compact_deleted_keys(jetstream.js, jetstream.bucket, older_than=0) + assert (await jetstream.js.stream_info(stream)).state.messages == 0 + + sent = jetstream.nc.stats["out_msgs"] + # Longer than the live watcher's inactive threshold. It must remain + # live while unused, and short-lived compaction consumers must expire. + await asyncio.sleep(35) + assert await store.claim_users("1", "idle", 100, 30) == [] + assert jetstream.nc.stats["out_msgs"] == sent + assert (await jetstream.js.stream_info(stream)).state.consumer_count == 1 + await store.enqueue_users("1", [User(email="after-idle")]) + assert [item.user.email for item in await store.claim_users("1", "worker", 100, 30)] == ["after-idle"] + finally: + await store.close() + + +async def test_live_index_survives_consumer_recreation(jetstream): + store = NatsUserSyncStore(jetstream.kv) + try: + assert await store.claim_users("1", "reader", 10, 30) == [] + consumers = await jetstream.js.consumers_info(f"KV_{jetstream.bucket}") + assert len(consumers) == 1 + await jetstream.js.delete_consumer(f"KV_{jetstream.bucket}", consumers[0].name) + writer = NatsUserSyncStore(jetstream.kv) + try: + await writer.enqueue_users("1", [User(email="after-consumer-loss")]) + finally: + await writer.close() + async with asyncio.timeout(30): + while not await store._key_index.keys("p.1."): + await asyncio.sleep(0.05) + assert [item.user.email for item in await store.claim_users("1", "reader", 10, 30)] == ["after-consumer-loss"] + finally: + await store.close() diff --git a/tests/test_record_usages.py b/tests/test_record_usages.py index 8a2b5ec2e..3397f7bf4 100644 --- a/tests/test_record_usages.py +++ b/tests/test_record_usages.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging import os from collections import defaultdict from typing import Any @@ -510,27 +511,89 @@ async def always_deadlock(_stmt, _params=None): await record_usages.safe_execute("stmt", [{"uid": 1}], max_retries=3) +@pytest.fixture(autouse=True) +def _reset_usage_job_state(): + record_usages._usage_coefficient_cache.clear() + record_usages._user_usage_running = False + record_usages._node_usage_running = False + yield + record_usages._usage_coefficient_cache.clear() + record_usages._user_usage_running = False + record_usages._node_usage_running = False + + @pytest.mark.asyncio -async def test_record_user_usages_skips_when_already_running(monkeypatch: pytest.MonkeyPatch): +async def test_record_user_usages_skips_when_already_running(monkeypatch: pytest.MonkeyPatch, caplog): impl = AsyncMock() monkeypatch.setattr(record_usages, "_record_user_usages_impl", impl) record_usages._user_usage_running = True - try: - await record_usages.record_user_usages() - finally: - record_usages._user_usage_running = False + caplog.set_level(logging.WARNING) + await record_usages.record_user_usages() impl.assert_not_awaited() + assert "JOB_RECORD_USER_USAGES_INTERVAL" in caplog.text + assert "UVICORN_WORKERS" in caplog.text @pytest.mark.asyncio -async def test_record_node_usages_skips_when_already_running(monkeypatch: pytest.MonkeyPatch): +async def test_record_node_usages_skips_when_already_running(monkeypatch: pytest.MonkeyPatch, caplog): impl = AsyncMock() monkeypatch.setattr(record_usages, "_record_node_usages_impl", impl) record_usages._node_usage_running = True - try: - await record_usages.record_node_usages() - finally: - record_usages._node_usage_running = False + caplog.set_level(logging.WARNING) + await record_usages.record_node_usages() impl.assert_not_awaited() + assert "JOB_RECORD_NODE_USAGES_INTERVAL" in caplog.text + assert "UVICORN_WORKERS" in caplog.text + + +@pytest.mark.asyncio +async def test_record_user_usages_does_not_apply_global_timeout(monkeypatch: pytest.MonkeyPatch): + impl = AsyncMock() + wait_for = AsyncMock(side_effect=AssertionError("wait_for should not run")) + monkeypatch.setattr(record_usages, "_record_user_usages_impl", impl) + monkeypatch.setattr(record_usages.asyncio, "wait_for", wait_for) + + await record_usages.record_user_usages() + + impl.assert_awaited_once() + wait_for.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_record_user_usages_warns_when_slower_than_interval(monkeypatch: pytest.MonkeyPatch, caplog): + clock = {"t": 0.0} + + async def impl(): + clock["t"] = 15.0 + + monkeypatch.setattr(record_usages.time, "monotonic", lambda: clock["t"]) + monkeypatch.setattr(record_usages, "_record_user_usages_impl", impl) + monkeypatch.setattr(record_usages.job_settings, "record_user_usages_interval", 10) + caplog.set_level(logging.WARNING) + + await record_usages.record_user_usages() + + assert "exceeds the 10s interval" in caplog.text + assert "UVICORN_WORKERS" in caplog.text + + +@pytest.mark.asyncio +async def test_usage_coefficient_is_cached_across_collects(monkeypatch: pytest.MonkeyPatch): + node = DummyNode(1, usage_coefficient=2) + extra_calls = {"n": 0} + original_get_extra = node.get_extra + + async def counting_get_extra(): + extra_calls["n"] += 1 + return await original_get_extra() + + node.get_extra = counting_get_extra + monkeypatch.setattr(record_usages, "get_users_stats", AsyncMock(return_value=[])) + + first = await record_usages._collect_node_user_usage(node, 1) + second = await record_usages._collect_node_user_usage(node, 1) + + assert extra_calls["n"] == 1 + assert first[1] == second[1] == 2.0 diff --git a/tests/test_review_users_unit.py b/tests/test_review_users_unit.py new file mode 100644 index 000000000..e7a192a99 --- /dev/null +++ b/tests/test_review_users_unit.py @@ -0,0 +1,175 @@ +"""Expire/limit review jobs sync users in bulk instead of per row.""" + +from __future__ import annotations + +from datetime import UTC, datetime as dt, timedelta as td +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest +from sqlalchemy import inspect as sa_inspect +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlalchemy.pool import StaticPool + +from app.db import base +from app.db.crud.user import get_active_to_expire_users, start_users_expire +from app.db.models import User, UserStatus, UserUsageResetLogs +from app.jobs import review_users +from app.models.settings import HWIDSettings +from app.operation import OperatorType +from app.operation.subscription import SubscriptionOperation + + +def _user(user_id: int, *, next_plan=None) -> SimpleNamespace: + return SimpleNamespace(id=user_id, username=f"u{user_id}", next_plan=next_plan) + + +@pytest.mark.asyncio +async def test_apply_status_changes_syncs_plain_users_in_one_call(): + users = [_user(1), _user(2), _user(3)] + + async def _validate(db_user, include_subscription_url=True): + return SimpleNamespace(id=db_user.id, username=db_user.username) + + with ( + patch("app.jobs.review_users.update_users_status", new_callable=AsyncMock) as update_status, + patch("app.jobs.review_users.sync_users", new_callable=AsyncMock) as sync, + patch.object(review_users.user_operator, "validate_user", side_effect=_validate), + patch("app.jobs.review_users.notification.user_status_change", new_callable=AsyncMock) as status_change, + patch("app.jobs.review_users.reset_user_by_next", new_callable=AsyncMock) as reset_next, + ): + await review_users.apply_status_changes(AsyncMock(), users, UserStatus.expired) + + update_status.assert_awaited_once() + assert update_status.await_args.args[1] == users + sync.assert_awaited_once_with(users) + reset_next.assert_not_called() + assert status_change.call_count == 3 + + +@pytest.mark.asyncio +async def test_apply_status_changes_resets_next_plan_users_separately(): + next_plan = SimpleNamespace(data_limit=1) + plain = _user(1) + planned = _user(2, next_plan=next_plan) + reset_user = _user(2) + reset_user.status = UserStatus.active + + async def _validate(db_user, include_subscription_url=True): + return SimpleNamespace(id=db_user.id, username=db_user.username) + + with ( + patch("app.jobs.review_users.update_users_status", new_callable=AsyncMock) as update_status, + patch("app.jobs.review_users.sync_users", new_callable=AsyncMock) as sync, + patch.object(review_users.user_operator, "validate_user", side_effect=_validate), + patch("app.jobs.review_users.notification.user_status_change", new_callable=AsyncMock), + patch("app.jobs.review_users.notification.user_data_reset_by_next", new_callable=AsyncMock) as reset_notify, + patch( + "app.jobs.review_users.reset_user_by_next", new_callable=AsyncMock, return_value=reset_user + ) as reset_next, + ): + await review_users.apply_status_changes(AsyncMock(), [plain, planned], UserStatus.limited) + + update_status.assert_awaited_once() + assert update_status.await_args.args[1] == [plain] + reset_next.assert_awaited_once() + assert [call.args[0] for call in sync.await_args_list] == [[plain], [reset_user]] + reset_notify.assert_called_once() + + +@pytest.mark.asyncio +async def test_apply_status_changes_skips_status_update_when_already_active(): + users = [_user(1)] + with ( + patch("app.jobs.review_users.update_users_status", new_callable=AsyncMock) as update_status, + patch("app.jobs.review_users.sync_users", new_callable=AsyncMock) as sync, + patch.object( + review_users.user_operator, + "validate_user", + new_callable=AsyncMock, + return_value=SimpleNamespace(id=1, username="u1"), + ), + patch("app.jobs.review_users.notification.user_status_change", new_callable=AsyncMock), + ): + await review_users.apply_status_changes(AsyncMock(), users, UserStatus.active) + + update_status.assert_not_called() + sync.assert_awaited_once_with(users) + + +@pytest.fixture +async def db_session(): + engine = create_async_engine( + "sqlite+aiosqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + async with engine.begin() as conn: + await conn.run_sync(base.Base.metadata.create_all) + factory = async_sessionmaker(bind=engine, expire_on_commit=False) + async with factory() as session: + yield session + await engine.dispose() + + +@pytest.mark.asyncio +async def test_expire_query_skips_usage_logs(db_session): + user = User(username="expiring", status=UserStatus.active) + user.expire = dt.now(UTC) - td(hours=1) + db_session.add(user) + await db_session.flush() + db_session.add(UserUsageResetLogs(user_id=user.id, used_traffic_at_reset=9)) + await db_session.commit() + + expired = await get_active_to_expire_users(db_session) + assert len(expired) == 1 + assert "usage_logs" in sa_inspect(expired[0]).unloaded + assert expired[0].lifetime_used_traffic == 9 + + +@pytest.mark.asyncio +async def test_start_users_expire_updates_all_rows_in_one_statement(db_session): + users = [ + User(username="hold1", status=UserStatus.on_hold, on_hold_expire_duration=3600), + User(username="hold2", status=UserStatus.on_hold, on_hold_expire_duration=7200), + ] + db_session.add_all(users) + await db_session.commit() + + updated = await start_users_expire(db_session, users) + assert all(user.status == UserStatus.active for user in updated) + assert all(user.on_hold_expire_duration is None for user in updated) + assert updated[0].expire is not None + assert updated[1].expire is not None + + +@pytest.mark.asyncio +async def test_existing_hwid_skips_register_when_recently_used(): + op = SubscriptionOperation(OperatorType.API) + existing = SimpleNamespace(last_used_at=dt.now(UTC) - td(seconds=30)) + settings = HWIDSettings(enabled=True, forced=False, fallback_limit=3) + + with ( + patch("app.operation.subscription.hwid_settings", new_callable=AsyncMock, return_value=settings), + patch("app.operation.subscription.get_user_hwid_by_value", new_callable=AsyncMock, return_value=existing), + patch("app.operation.subscription.register_user_hwid", new_callable=AsyncMock) as register, + ): + await op.validate_and_register_hwid(AsyncMock(), 1, None, None, "device-1", None, None, None) + + register.assert_not_called() + + +@pytest.mark.asyncio +async def test_existing_hwid_registers_when_stale(): + op = SubscriptionOperation(OperatorType.API) + existing = SimpleNamespace(last_used_at=dt.now(UTC) - td(minutes=10)) + settings = HWIDSettings(enabled=True, forced=False, fallback_limit=3) + + with ( + patch("app.operation.subscription.hwid_settings", new_callable=AsyncMock, return_value=settings), + patch("app.operation.subscription.get_user_hwid_by_value", new_callable=AsyncMock, return_value=existing), + patch("app.operation.subscription.register_user_hwid", new_callable=AsyncMock) as register, + ): + await op.validate_and_register_hwid(AsyncMock(), 1, None, None, "device-1", "iOS", "16", "iPhone") + + register.assert_awaited_once() diff --git a/tests/test_sub_update_buffer.py b/tests/test_sub_update_buffer.py new file mode 100644 index 000000000..b092f1fbe --- /dev/null +++ b/tests/test_sub_update_buffer.py @@ -0,0 +1,202 @@ +"""Buffered user_subscription_updates writes stay off the request commit path.""" + +from __future__ import annotations + +import asyncio + +import pytest +from sqlalchemy import delete, event, func, select +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlalchemy.pool import StaticPool + +from app.db import base +from app.db.crud.user import get_users_sub_update_list, user_sub_update +from app.db.models import User, UserSubscriptionUpdate +from app.subscription import sub_update_buffer + + +@pytest.fixture +async def buffer_db(monkeypatch: pytest.MonkeyPatch): + engine = create_async_engine( + "sqlite+aiosqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + + @event.listens_for(engine.sync_engine, "connect") + def enable_foreign_keys(connection, _record): + cursor = connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + + async with engine.begin() as conn: + await conn.run_sync(base.Base.metadata.create_all) + + factory = async_sessionmaker(bind=engine, expire_on_commit=False, autoflush=False) + + class TestGetDB: + def __init__(self): + self.db = factory() + + async def __aenter__(self): + return self.db + + async def __aexit__(self, exc_type, exc_value, traceback): + try: + if exc_type is not None: + await self.db.rollback() + finally: + await self.db.close() + + monkeypatch.setattr(sub_update_buffer, "GetDB", TestGetDB) + await sub_update_buffer.reset_user_sub_update_buffer() + + async with factory() as session: + session.add(User(username="buffered")) + await session.commit() + user_id = (await session.execute(select(User.id).where(User.username == "buffered"))).scalar_one() + yield session, user_id + + await sub_update_buffer.reset_user_sub_update_buffer() + await engine.dispose() + + +@pytest.mark.asyncio +async def test_queue_does_not_write_until_flush(buffer_db): + session, user_id = buffer_db + await sub_update_buffer.queue_user_sub_update(user_id, "v2rayNG/1.0", ip="203.0.113.10", hwid="abc") + + count = (await session.execute(select(func.count()).select_from(UserSubscriptionUpdate))).scalar() + assert count == 0 + assert sub_update_buffer.pending_count() == 1 + await session.commit() + + written = await sub_update_buffer.flush_user_sub_updates() + assert written == 1 + assert sub_update_buffer.pending_count() == 0 + + rows = (await session.execute(select(UserSubscriptionUpdate))).scalars().all() + assert len(rows) == 1 + assert rows[0].user_agent == "v2rayNG/1.0" + assert rows[0].ip == "203.0.113.10" + assert rows[0].hwid == "abc" + + +@pytest.mark.asyncio +async def test_user_sub_update_truncates_and_list_flushes(buffer_db): + session, user_id = buffer_db + await user_sub_update(session, user_id, "A" * 1000, ip="1.2.3.4") + assert sub_update_buffer.pending_count() == 1 + + stored, count = await get_users_sub_update_list(session, user_id) + assert count == 1 + assert stored[0].user_agent == "A" * 512 + assert sub_update_buffer.pending_count() == 0 + + +@pytest.mark.asyncio +async def test_flush_failure_requeues(buffer_db, monkeypatch: pytest.MonkeyPatch): + _session, user_id = buffer_db + await sub_update_buffer.queue_user_sub_update(user_id, "clash") + + class BoomGetDB: + def __init__(self): + raise SQLAlchemyError("boom") + + monkeypatch.setattr(sub_update_buffer, "GetDB", BoomGetDB) + with pytest.raises(SQLAlchemyError): + await sub_update_buffer.flush_user_sub_updates() + assert sub_update_buffer.pending_count() == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("include_live_user", [False, True]) +@pytest.mark.parametrize("batch_size", [1, 100]) +async def test_flush_discards_deleted_users(buffer_db, monkeypatch, include_live_user, batch_size): + session, user_id = buffer_db + deleted_user = User(username="deleted_before_flush") + session.add(deleted_user) + await session.commit() + await sub_update_buffer.queue_user_sub_update(deleted_user.id, "deleted-client") + if include_live_user: + await sub_update_buffer.queue_user_sub_update(user_id, "live-client") + + await session.execute(delete(User).where(User.id == deleted_user.id)) + await session.commit() + monkeypatch.setattr(sub_update_buffer, "FLUSH_BATCH_SIZE", batch_size) + + assert await sub_update_buffer.flush_user_sub_updates() == int(include_live_user) + assert sub_update_buffer.pending_count() == 0 + rows = (await session.execute(select(UserSubscriptionUpdate))).scalars().all() + assert [(row.user_id, row.user_agent) for row in rows] == ([(user_id, "live-client")] if include_live_user else []) + await session.commit() + + monkeypatch.setattr(sub_update_buffer, "FLUSH_BATCH_SIZE", 100) + await sub_update_buffer.queue_user_sub_update(user_id, "next-client") + assert await sub_update_buffer.flush_user_sub_updates() == 1 + assert sub_update_buffer.pending_count() == 0 + + +@pytest.mark.asyncio +async def test_second_flush_waits_until_in_flight_drain_commits(buffer_db, monkeypatch): + session, user_id = buffer_db + inner_cls = sub_update_buffer.GetDB + first_entered = asyncio.Event() + release = asyncio.Event() + calls = {"n": 0} + + class WrappedGetDB: + def __init__(self): + self._inner = inner_cls() + + async def __aenter__(self): + calls["n"] += 1 + if calls["n"] == 1: + first_entered.set() + await release.wait() + return await self._inner.__aenter__() + + async def __aexit__(self, exc_type, exc_value, traceback): + return await self._inner.__aexit__(exc_type, exc_value, traceback) + + monkeypatch.setattr(sub_update_buffer, "GetDB", WrappedGetDB) + await sub_update_buffer.queue_user_sub_update(user_id, "first-client") + first_flush = asyncio.create_task(sub_update_buffer.flush_user_sub_updates()) + await first_entered.wait() + await sub_update_buffer.queue_user_sub_update(user_id, "second-client") + second_flush = asyncio.create_task(sub_update_buffer.flush_user_sub_updates()) + await asyncio.sleep(0.05) + assert not second_flush.done() + release.set() + written = await first_flush + await second_flush + assert written == 2 + assert second_flush.done() + assert sub_update_buffer.pending_count() == 0 + rows = (await session.execute(select(UserSubscriptionUpdate))).scalars().all() + assert sorted(row.user_agent for row in rows) == ["first-client", "second-client"] + await session.commit() + + +@pytest.mark.asyncio +async def test_queue_spawns_flush_only_when_crossing_batch_size(buffer_db, monkeypatch): + _session, user_id = buffer_db + created: list[str | None] = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro, *args, **kwargs): + created.append(kwargs.get("name")) + coro.close() + return real_create_task(asyncio.sleep(0), name=kwargs.get("name")) + + monkeypatch.setattr(sub_update_buffer.asyncio, "create_task", tracking_create_task) + monkeypatch.setattr(sub_update_buffer, "FLUSH_BATCH_SIZE", 2) + + await sub_update_buffer.queue_user_sub_update(user_id, "one") + assert created == [] + await sub_update_buffer.queue_user_sub_update(user_id, "two") + assert created == ["sub_update_flush"] + await sub_update_buffer.queue_user_sub_update(user_id, "three") + assert created == ["sub_update_flush"] + await sub_update_buffer.reset_user_sub_update_buffer() + assert sub_update_buffer.pending_count() == 0 diff --git a/tests/test_subscription_cpu.py b/tests/test_subscription_cpu.py new file mode 100644 index 000000000..c4164ca65 --- /dev/null +++ b/tests/test_subscription_cpu.py @@ -0,0 +1,142 @@ +"""CPU-oriented subscription render, cache, and host-copy behaviour.""" + +from __future__ import annotations + +import json +import time +from collections import defaultdict +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from app.db.models import UserStatus +from app.models.subscription import SubscriptionInboundData, TCPTransportConfig, TLSConfig +from app.subscription import config_cache +from app.subscription.base import dumps_compact +from app.subscription.clash import ClashConfiguration +from app.subscription.outline import OutlineConfiguration +from app.subscription.share import generate_subscription, process_host +from app.subscription.singbox import SingBoxConfiguration +from app.subscription.xray import XrayConfiguration + + +@pytest.fixture(autouse=True) +def _clear_sub_config_cache(): + config_cache.clear_sub_config_cache() + yield + config_cache.clear_sub_config_cache() + + +def test_dumps_compact_has_no_pretty_whitespace(): + payload = [{"remarks": "a", "nested": {"k": 1}}] + rendered = dumps_compact(payload) + assert "\n" not in rendered + assert " " not in rendered + assert json.loads(rendered) == payload + + +def test_xray_and_singbox_render_compact(): + xray = XrayConfiguration() + xray.config = [{"remarks": "node", "outbounds": [{"tag": "direct"}]}] + xray_text = xray.render() + assert "\n" not in xray_text + assert json.loads(xray_text) == xray.config + + singbox = SingBoxConfiguration(singbox_template_content="{}") + singbox.config = {"outbounds": [{"type": "direct", "tag": "direct"}], "endpoints": []} + singbox_text = singbox.render() + assert "\n" not in singbox_text + parsed = json.loads(singbox_text) + assert parsed["outbounds"][0]["tag"] == "direct" + + +def test_outline_render_compact(): + outline = OutlineConfiguration() + outline.add_directly({"method": "aes-256-gcm", "password": "x"}) + text = outline.render() + assert "\n" not in text + assert json.loads(text) == {"method": "aes-256-gcm", "password": "x"} + + +def test_clash_render_returns_template_output_without_yaml_roundtrip(): + conf = ClashConfiguration(clash_template_content="count={{ proxy_remarks | length }}") + conf.proxy_remarks = ["a", "b"] + assert conf.render() == "count=2" + + +def test_sub_config_cache_expires_and_evicts(): + key = (1, "links", False, False, "active", None, None, ()) + config_cache.put_sub_config(key, "v1") + assert config_cache.get_sub_config(key) == "v1" + + original_ttl = config_cache.SUB_CONFIG_CACHE_TTL_S + config_cache.SUB_CONFIG_CACHE_TTL_S = 0 + try: + config_cache.put_sub_config(key, "v2") + time.sleep(0.01) + assert config_cache.get_sub_config(key) is None + finally: + config_cache.SUB_CONFIG_CACHE_TTL_S = original_ttl + + original_max = config_cache.SUB_CONFIG_CACHE_MAX + config_cache.SUB_CONFIG_CACHE_MAX = 2 + try: + config_cache.clear_sub_config_cache() + config_cache.put_sub_config(("a",), "1") + config_cache.put_sub_config(("b",), "2") + config_cache.put_sub_config(("c",), "3") + assert config_cache.get_sub_config(("a",)) is None + assert config_cache.get_sub_config(("b",)) == "2" + assert config_cache.get_sub_config(("c",)) == "3" + finally: + config_cache.SUB_CONFIG_CACHE_MAX = original_max + + +@pytest.mark.asyncio +async def test_generate_subscription_returns_cached_payload(monkeypatch: pytest.MonkeyPatch): + user = SimpleNamespace(id=7, status=UserStatus.active, inbounds=["tag"], data_limit=None, expire=None) + monkeypatch.setattr( + "app.subscription.share.subscription_client_templates", + AsyncMock(side_effect=AssertionError("cache hit should skip rebuild")), + ) + key = config_cache.make_sub_config_key(user, "links", False, False) + config_cache.put_sub_config(key, "cached-links") + assert await generate_subscription(user, "links", False) == "cached-links" + + +@pytest.mark.asyncio +async def test_process_host_does_not_mutate_cached_inbound(): + inbound = SubscriptionInboundData( + remark="{USERNAME}", + inbound_tag="in1", + protocol="vless", + address=["host.example"], + port=[443], + network="tcp", + tls_config=TLSConfig(tls="tls", sni=["sni.example"], reality_short_id="abc"), + transport_config=TCPTransportConfig(path="/{USERNAME}", host=["cdn.example"]), + priority=0, + ) + original_sni = list(inbound.tls_config.sni) + original_host = list(inbound.transport_config.host) + original_path = inbound.transport_config.path + + result = await process_host( + inbound, + defaultdict(lambda: "", {"USERNAME": "u1"}), + ["in1"], + {"vless": {"id": "11111111-1111-1111-1111-111111111111"}}, + ) + + assert result is not None + copy, settings = result + assert inbound.tls_config.sni == original_sni + assert inbound.transport_config.host == original_host + assert inbound.transport_config.path == original_path + assert copy is not inbound + assert copy.tls_config is not inbound.tls_config + assert copy.transport_config is not inbound.transport_config + assert copy.tls_config.sni == "sni.example" + assert copy.transport_config.path == "/u1" + assert settings["id"] == "11111111-1111-1111-1111-111111111111" diff --git a/tests/test_subscription_hot_path.py b/tests/test_subscription_hot_path.py new file mode 100644 index 000000000..a6a3265a9 --- /dev/null +++ b/tests/test_subscription_hot_path.py @@ -0,0 +1,104 @@ +"""Slim GET /sub user load and inbounds() reuse of already-loaded groups.""" + +from __future__ import annotations + +import pytest +from sqlalchemy import event, inspect as sa_inspect, select +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlalchemy.orm import selectinload +from sqlalchemy.pool import StaticPool + +from app.db import base +from app.db.crud.user import get_user_by_id +from app.db.models import Group, ProxyInbound, User, UserUsageResetLogs, users_groups_association +from app.operation.subscription import SubscriptionOperation + + +@pytest.fixture +async def db_session(): + engine = create_async_engine( + "sqlite+aiosqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + async with engine.begin() as conn: + await conn.run_sync(base.Base.metadata.create_all) + + factory = async_sessionmaker(bind=engine, expire_on_commit=False) + async with factory() as session: + inbound = ProxyInbound(tag="vless-tcp") + group = Group(name="g1", inbounds=[inbound]) + user = User(username="subuser", used_traffic=1000) + session.add_all([inbound, group, user]) + await session.flush() + await session.execute(users_groups_association.insert().values(user_id=user.id, groups_id=group.id)) + session.add(UserUsageResetLogs(user_id=user.id, used_traffic_at_reset=234)) + await session.commit() + yield session, user.id, engine + + await engine.dispose() + + +@pytest.mark.asyncio +async def test_slim_user_fetch_skips_heavy_relations(db_session): + session, user_id, _engine = db_session + user = await get_user_by_id( + session, + user_id, + load_admin=True, + load_admin_role=True, + load_next_plan=False, + load_usage_logs=False, + load_groups=False, + load_lifetime_used_traffic=True, + ) + assert user is not None + unloaded = sa_inspect(user).unloaded + assert "usage_logs" in unloaded + assert "next_plan" in unloaded + assert "groups" in unloaded + assert user.lifetime_used_traffic == 1234 + assert "vless-tcp" in await user.inbounds() + assert "groups" in sa_inspect(user).unloaded + + +@pytest.mark.asyncio +async def test_inbounds_skips_sql_when_groups_and_inbounds_are_loaded(db_session): + session, user_id, engine = db_session + user = ( + await session.execute( + select(User).options(selectinload(User.groups).selectinload(Group.inbounds)).where(User.id == user_id) + ) + ).scalar_one() + + select_count = 0 + + def _count(_conn, _cursor, statement, *_args): + nonlocal select_count + if str(statement).lstrip().lower().startswith("select"): + select_count += 1 + + event.listen(engine.sync_engine, "before_cursor_execute", _count) + try: + tags = await user.inbounds() + finally: + event.remove(engine.sync_engine, "before_cursor_execute", _count) + + assert tags == ["vless-tcp"] + assert select_count == 0 + + +@pytest.mark.asyncio +async def test_validated_user_uses_lifetime_expression(db_session): + session, user_id, _engine = db_session + db_user = await get_user_by_id( + session, + user_id, + load_next_plan=False, + load_usage_logs=False, + load_groups=False, + load_lifetime_used_traffic=True, + ) + user = await SubscriptionOperation.validated_user(db_user) + assert user.lifetime_used_traffic == 1234 + assert user.inbounds == ["vless-tcp"]