From 85e5f3e9a44e107f5b27c7daaa3addd209d0af0a Mon Sep 17 00:00:00 2001 From: x0sina Date: Wed, 9 Sep 2026 17:36:27 +0330 Subject: [PATCH 1/6] feat(user): enhance user inbound tag retrieval and subscription update handling - Improved the `inbounds` method in the User model to efficiently gather inbound tags from enabled groups. - Updated user subscription update logic to queue updates for background processing, enhancing performance. - Refactored user status change notifications to streamline the process and ensure accurate user state management. - Introduced caching for usage coefficients to optimize node user statistics collection. - Enhanced subscription generation with caching mechanisms for improved performance. --- app/db/crud/user.py | 80 +++--- app/db/models.py | 17 +- app/jobs/cleanup_subscription_updates.py | 3 + app/jobs/record_usages.py | 316 +++++++++-------------- app/jobs/reset_user_data_usage.py | 4 +- app/jobs/review_users.py | 69 +++-- app/operation/__init__.py | 25 +- app/operation/subscription.py | 42 ++- app/subscription/base.py | 5 + app/subscription/clash.py | 16 +- app/subscription/config_cache.py | 63 +++++ app/subscription/outline.py | 6 +- app/subscription/share.py | 51 ++-- app/subscription/singbox.py | 3 +- app/subscription/sub_update_buffer.py | 144 +++++++++++ app/subscription/xray.py | 3 +- app/templates/filters.py | 5 +- tests/api/conftest.py | 1 + tests/test_record_usages.py | 83 +++++- tests/test_review_users_unit.py | 175 +++++++++++++ tests/test_sub_update_buffer.py | 101 ++++++++ tests/test_subscription_cpu.py | 142 ++++++++++ tests/test_subscription_hot_path.py | 104 ++++++++ 23 files changed, 1148 insertions(+), 310 deletions(-) create mode 100644 app/subscription/config_cache.py create mode 100644 app/subscription/sub_update_buffer.py create mode 100644 tests/test_review_users_unit.py create mode 100644 tests/test_sub_update_buffer.py create mode 100644 tests/test_subscription_cpu.py create mode 100644 tests/test_subscription_hot_path.py diff --git a/app/db/crud/user.py b/app/db/crud/user.py index 46e5def36..9e16cf914 100644 --- a/app/db/crud/user.py +++ b/app/db/crud/user.py @@ -3,7 +3,7 @@ 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 +47,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 +66,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 +603,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 +648,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 +677,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 +704,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 +712,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,33 +1348,32 @@ async def bulk_revoke_user_sub( return users +async def _flush_pending_subscription_updates() -> None: + await flush_user_sub_updates() + + 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( db: AsyncSession, user_id: int, offset: int = 0, limit: int = 10 ) -> tuple[Sequence[UserSubscriptionUpdate], int]: + await _flush_pending_subscription_updates() stmt = ( select(UserSubscriptionUpdate) .where(UserSubscriptionUpdate.user_id == user_id) @@ -1408,6 +1416,7 @@ async def get_users_subscription_agent_counts( end: datetime | None = None, period: Period | None = None, ) -> list[tuple[str, int]]: + await _flush_pending_subscription_updates() stmt = select(UserSubscriptionUpdate.user_agent, func.count().label("count")) from_clause, conditions = _subscription_update_from_clause(user_id=user_id, admin_id=admin_id) @@ -1437,6 +1446,7 @@ async def get_users_subscription_agent_stats( admin_id: int | None = None, ) -> list[dict]: """Retrieve subscription update counts grouped by agent and period.""" + await _flush_pending_subscription_updates() trunc_expr = _build_trunc_expression(db, period, UserSubscriptionUpdate.created_at, start) start_utc = get_complete_period_start_for_filter(start, period) end_utc = to_utc_for_filter(end) @@ -1825,7 +1835,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 +1847,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 30ed85024..78507474a 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_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/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..f452fb8b7 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,8 @@ 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.config_cache import SUB_CONFIG_CACHE_TTL_S from app.subscription.share import ( apply_custom_format_variables, encode_title, @@ -87,6 +90,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: @@ -197,6 +212,7 @@ def create_response_headers( "subscription-userinfo": "; ".join(f"{key}={val}" for key, val in user_info.items()), "announce": encode_title(formatted_announce), "announce-url": formatted_announce_url, + "Cache-Control": f"private, max-age={SUB_CONFIG_CACHE_TTL_S}", } if extra_headers: headers.update(extra_headers) @@ -389,6 +405,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 +438,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 +561,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 +633,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 +703,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 +720,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 +759,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 +812,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 f96c1ac9b..b2faba5e8 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 409c0e32d..082a52746 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..8a0abe545 --- /dev/null +++ b/app/subscription/sub_update_buffer.py @@ -0,0 +1,144 @@ +"""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 + +from app.db import GetDB +from app.db.models import 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() +_flushing = False +_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.""" + global _flushing + async with _lock: + _pending.clear() + _flushing = False + + +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: + _pending.append(record) + overflow = len(_pending) - _MAX_BUFFER + if overflow > 0: + del _pending[:overflow] + dropped = overflow + should_flush = len(_pending) >= FLUSH_BATCH_SIZE + 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. Returns the number of rows written in this call.""" + global _flushing + written = 0 + while True: + async with _lock: + if _flushing: + return written + if not _pending: + return written + _flushing = True + batch = _pending[:] + _pending.clear() + try: + async with GetDB() as db: + await db.execute(insert(UserSubscriptionUpdate.__table__), batch) + await db.commit() + written += len(batch) + 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 + finally: + async with _lock: + _flushing = False + + +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/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/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..3a4a7499c --- /dev/null +++ b/tests/test_sub_update_buffer.py @@ -0,0 +1,101 @@ +"""Buffered user_subscription_updates writes stay off the request commit path.""" + +from __future__ import annotations + +import pytest +from sqlalchemy import 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, + ) + 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 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"] From ce739e2f406fecb1bc0b14f1e7c26a8579243a09 Mon Sep 17 00:00:00 2001 From: DrSaeedHub <122322581+DrSaeedHub@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:38:22 +0000 Subject: [PATCH 2/6] fix: discard stale subscription updates and align FinalMask test (#886) --- app/subscription/sub_update_buffer.py | 26 ++++++++++++++----- tests/api/test_host.py | 5 +++- tests/test_sub_update_buffer.py | 37 ++++++++++++++++++++++++++- 3 files changed, 60 insertions(+), 8 deletions(-) diff --git a/app/subscription/sub_update_buffer.py b/app/subscription/sub_update_buffer.py index 8a0abe545..e5fa83abe 100644 --- a/app/subscription/sub_update_buffer.py +++ b/app/subscription/sub_update_buffer.py @@ -10,10 +10,10 @@ from datetime import UTC, datetime as dt from typing import Any -from sqlalchemy import insert +from sqlalchemy import insert, select from app.db import GetDB -from app.db.models import UserSubscriptionUpdate +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 @@ -87,13 +87,27 @@ async def flush_user_sub_updates() -> int: if not _pending: return written _flushing = True - batch = _pending[:] - _pending.clear() + batch = _pending[:FLUSH_BATCH_SIZE] + del _pending[:FLUSH_BATCH_SIZE] try: async with GetDB() as db: - await db.execute(insert(UserSubscriptionUpdate.__table__), batch) + # 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(batch) + written += len(live_records) except Exception: async with _lock: _pending[0:0] = batch diff --git a/tests/api/test_host.py b/tests/api/test_host.py index 1bc21a3a4..2bc8a9a8b 100644 --- a/tests/api/test_host.py +++ b/tests/api/test_host.py @@ -344,7 +344,7 @@ def test_host_finalmask_new_types(access_token): "type": "noise", "settings": { "reset": "30-60", - "noise": [{"type": "rand", "rand": "1-8192", "delay": "10-20"}], + "noise": [{"type": "array", "packet": [1, 2, 3], "rand": "1-8192", "delay": "10-20"}], }, }, ], @@ -389,6 +389,9 @@ def test_host_finalmask_new_types(access_token): assert fm["udp"][0]["type"] == "realm" assert fm["udp"][1]["type"] == "mkcp-legacy" assert fm["udp"][5]["settings"].get("reset") == "30-60" + assert fm["udp"][5]["settings"]["noise"][0]["type"] == "array" + assert fm["udp"][5]["settings"]["noise"][0]["packet"] == [1, 2, 3] + assert fm["udp"][5]["settings"]["noise"][0]["rand"] == "1-8192" assert "apply_to" not in (fm["udp"][5]["settings"].get("noise") or [{}])[0] finally: client.delete(f"/api/host/{host_id}", headers={"Authorization": f"Bearer {access_token}"}) diff --git a/tests/test_sub_update_buffer.py b/tests/test_sub_update_buffer.py index 3a4a7499c..a35123694 100644 --- a/tests/test_sub_update_buffer.py +++ b/tests/test_sub_update_buffer.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest -from sqlalchemy import func, select +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 @@ -21,6 +21,13 @@ async def buffer_db(monkeypatch: pytest.MonkeyPatch): 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) @@ -99,3 +106,31 @@ def __init__(self): 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 From 327afccb4c36d6bac606b140a428964f5bd17248 Mon Sep 17 00:00:00 2001 From: DrSaeedHub <122322581+DrSaeedHub@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:53:11 +0000 Subject: [PATCH 3/6] fix: read flushed subscription updates from a fresh snapshot (#888) --- app/db/crud/user.py | 36 +++++++---- .../api/test_subscription_update_snapshot.py | 64 +++++++++++++++++++ 2 files changed, 86 insertions(+), 14 deletions(-) create mode 100644 tests/api/test_subscription_update_snapshot.py diff --git a/app/db/crud/user.py b/app/db/crud/user.py index 9e16cf914..77d146322 100644 --- a/app/db/crud/user.py +++ b/app/db/crud/user.py @@ -1,4 +1,5 @@ -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 @@ -1348,8 +1349,15 @@ async def bulk_revoke_user_sub( return users -async def _flush_pending_subscription_updates() -> None: +@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( @@ -1373,22 +1381,22 @@ async def user_sub_update( async def get_users_sub_update_list( db: AsyncSession, user_id: int, offset: int = 0, limit: int = 10 ) -> tuple[Sequence[UserSubscriptionUpdate], int]: - await _flush_pending_subscription_updates() stmt = ( select(UserSubscriptionUpdate) .where(UserSubscriptionUpdate.user_id == user_id) .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 @@ -1416,7 +1424,6 @@ async def get_users_subscription_agent_counts( end: datetime | None = None, period: Period | None = None, ) -> list[tuple[str, int]]: - await _flush_pending_subscription_updates() stmt = select(UserSubscriptionUpdate.user_agent, func.count().label("count")) from_clause, conditions = _subscription_update_from_clause(user_id=user_id, admin_id=admin_id) @@ -1433,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()] @@ -1446,7 +1454,6 @@ async def get_users_subscription_agent_stats( admin_id: int | None = None, ) -> list[dict]: """Retrieve subscription update counts grouped by agent and period.""" - await _flush_pending_subscription_updates() trunc_expr = _build_trunc_expression(db, period, UserSubscriptionUpdate.created_at, start) start_utc = get_complete_period_start_for_filter(start, period) end_utc = to_utc_for_filter(end) @@ -1470,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(): 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() From 03c8b5fc723bc47de55c92fe518783636dc50936 Mon Sep 17 00:00:00 2001 From: DrSaeedHub <122322581+DrSaeedHub@users.noreply.github.com> Date: Sat, 12 Sep 2026 06:48:25 +0000 Subject: [PATCH 4/6] perf(nats): optimize user sync and bound queue processing (#892) Reuse live key indexes, replay snapshots in bounded batches, and clear queues from cached keys and recent revisions. Bound bulk operations, safely compact old deletion markers, and preserve claim recovery across delayed notifications and reconnects. Add concurrency, recovery, compaction, and idle-traffic regression coverage. --- app/jobs/cleanup_node_sync.py | 37 ++ app/nats/kv_cas.py | 14 +- app/nats/kv_cleanup.py | 35 ++ app/nats/kv_index.py | 140 ++++++++ app/nats/kv_watch.py | 159 +++++++++ app/node/nats_memory.py | 75 ++++- tests/api/__init__.py | 18 +- tests/nats_sync_process_worker.py | 40 +++ tests/test_nats_kv_index.py | 194 +++++++++++ tests/test_nats_node_memory.py | 25 ++ tests/test_nats_sync_integration.py | 503 ++++++++++++++++++++++++++++ 11 files changed, 1221 insertions(+), 19 deletions(-) create mode 100644 app/jobs/cleanup_node_sync.py create mode 100644 app/nats/kv_cleanup.py create mode 100644 app/nats/kv_index.py create mode 100644 app/nats/kv_watch.py create mode 100644 tests/nats_sync_process_worker.py create mode 100644 tests/test_nats_kv_index.py create mode 100644 tests/test_nats_sync_integration.py diff --git a/app/jobs/cleanup_node_sync.py b/app/jobs/cleanup_node_sync.py new file mode 100644 index 000000000..ca9d468b0 --- /dev/null +++ b/app/jobs/cleanup_node_sync.py @@ -0,0 +1,37 @@ +"""Compact completed node-sync work on the scheduler leader only.""" + +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: + 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) + 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/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/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/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() From bb9c3486285bd7bc19ee75901786669a148a1d5a Mon Sep 17 00:00:00 2001 From: x0sina Date: Sat, 12 Sep 2026 10:48:20 +0330 Subject: [PATCH 5/6] test(hosts): match FinalMask noise fixture to Xray array packet Keep the #889 payload ([1, 2, 255], no rand) so it agrees with the merged assertions after pref-boost picked up both that change and the older #886 sample. --- tests/api/test_host.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/api/test_host.py b/tests/api/test_host.py index 0c65828e5..f972f2c1a 100644 --- a/tests/api/test_host.py +++ b/tests/api/test_host.py @@ -344,7 +344,13 @@ def test_host_finalmask_new_types(access_token): "type": "noise", "settings": { "reset": "30-60", - "noise": [{"type": "array", "packet": [1, 2, 3], "rand": "1-8192", "delay": "10-20"}], + "noise": [ + { + "type": "array", + "packet": [1, 2, 255], + "delay": "10-20", + } + ], }, }, ], From fbd57c75fae02e6c5610bffc4ba3758e7810181f Mon Sep 17 00:00:00 2001 From: x0sina Date: Sat, 12 Sep 2026 11:22:19 +0330 Subject: [PATCH 6/6] fix(cleanup): add timeout handling for node-sync key compaction and improve logging - Introduced asyncio timeout for the compaction process to prevent indefinite blocking. - Added a warning log for timeout occurrences to ensure visibility of potential issues. - Updated Cache-Control header in subscription responses to "no-store" for better caching behavior. - Refactored subscription update buffer to use a drain lock for concurrent flush operations, improving reliability and performance. --- app/jobs/cleanup_node_sync.py | 7 ++- app/operation/subscription.py | 3 +- app/subscription/sub_update_buffer.py | 83 ++++++++++++--------------- tests/test_sub_update_buffer.py | 66 +++++++++++++++++++++ 4 files changed, 111 insertions(+), 48 deletions(-) diff --git a/app/jobs/cleanup_node_sync.py b/app/jobs/cleanup_node_sync.py index ca9d468b0..6757401ae 100644 --- a/app/jobs/cleanup_node_sync.py +++ b/app/jobs/cleanup_node_sync.py @@ -1,5 +1,7 @@ """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 @@ -18,9 +20,12 @@ async def cleanup_node_sync(): if nc is None: return try: - count = await compact_deleted_keys(await get_jetstream_context(nc), nats_settings.node_user_sync_kv_bucket) + 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() diff --git a/app/operation/subscription.py b/app/operation/subscription.py index f452fb8b7..c5535ee7f 100644 --- a/app/operation/subscription.py +++ b/app/operation/subscription.py @@ -21,7 +21,6 @@ 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.config_cache import SUB_CONFIG_CACHE_TTL_S from app.subscription.share import ( apply_custom_format_variables, encode_title, @@ -212,10 +211,10 @@ def create_response_headers( "subscription-userinfo": "; ".join(f"{key}={val}" for key, val in user_info.items()), "announce": encode_title(formatted_announce), "announce-url": formatted_announce_url, - "Cache-Control": f"private, max-age={SUB_CONFIG_CACHE_TTL_S}", } if extra_headers: headers.update(extra_headers) + headers["Cache-Control"] = "no-store" return headers @classmethod diff --git a/app/subscription/sub_update_buffer.py b/app/subscription/sub_update_buffer.py index e5fa83abe..28379bde8 100644 --- a/app/subscription/sub_update_buffer.py +++ b/app/subscription/sub_update_buffer.py @@ -30,7 +30,7 @@ _pending: list[dict[str, Any]] = [] _lock = asyncio.Lock() -_flushing = False +_drain_lock = asyncio.Lock() _flush_task: asyncio.Task | None = None @@ -52,10 +52,8 @@ def pending_count() -> int: async def reset_user_sub_update_buffer() -> None: """Drop queued rows without writing. Tests only.""" - global _flushing async with _lock: _pending.clear() - _flushing = False async def queue_user_sub_update(user_id: int, user_agent: str, ip: str | None = None, hwid: str | None = None) -> None: @@ -64,12 +62,13 @@ async def queue_user_sub_update(user_id: int, user_agent: str, ip: str | None = 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 = len(_pending) >= FLUSH_BATCH_SIZE + should_flush = prev < FLUSH_BATCH_SIZE <= len(_pending) if dropped: logger.warning("Dropped %s buffered subscription updates; buffer full", dropped) if should_flush: @@ -77,48 +76,42 @@ async def queue_user_sub_update(user_id: int, user_agent: str, ip: str | None = async def flush_user_sub_updates() -> int: - """Persist queued rows. Returns the number of rows written in this call.""" - global _flushing - written = 0 - while True: - async with _lock: - if _flushing: - return written - if not _pending: - return written - _flushing = True - 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: + """Persist queued rows. Concurrent callers wait, then drain anything left.""" + async with _drain_lock: + written = 0 + while True: 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 - finally: - async with _lock: - _flushing = False + 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: diff --git a/tests/test_sub_update_buffer.py b/tests/test_sub_update_buffer.py index a35123694..b092f1fbe 100644 --- a/tests/test_sub_update_buffer.py +++ b/tests/test_sub_update_buffer.py @@ -2,6 +2,8 @@ from __future__ import annotations +import asyncio + import pytest from sqlalchemy import delete, event, func, select from sqlalchemy.exc import SQLAlchemyError @@ -134,3 +136,67 @@ async def test_flush_discards_deleted_users(buffer_db, monkeypatch, include_live 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