diff --git a/.env.example b/.env.example index 8b5427100..044dea138 100644 --- a/.env.example +++ b/.env.example @@ -47,7 +47,8 @@ UVICORN_PORT = 8000 ## Database pool settings are per worker, not global. Four workers with pool size 10 can open up to 40 base connections. # SQLALCHEMY_POOL_SIZE = 25 # SQLALCHEMY_MAX_OVERFLOW = 60 -# SQLALCHEMY_POOL_RECYCLE = 300 +# SQLALCHEMY_POOL_RECYCLE = 1800 +# SQLALCHEMY_POOL_TIMEOUT = 15 # ECHO_SQL_QUERIES = False ## NATS connection and subjects for multi-process/node coordination. diff --git a/app/core/xray.py b/app/core/xray.py index 7a85c9251..4b718fe1d 100644 --- a/app/core/xray.py +++ b/app/core/xray.py @@ -372,7 +372,7 @@ def _hysteria_finalmask_from_stream(stream: dict, net_settings: dict) -> dict | """Normalize Hysteria Salamander masks into finalmask for client generation.""" finalmask = stream.get("finalmask") or stream.get("finalMask") if isinstance(finalmask, dict): - finalmask = deepcopy(finalmask) + finalmask = {k: v for k, v in finalmask.items()} else: finalmask = {} @@ -383,7 +383,7 @@ def _hysteria_finalmask_from_stream(stream: dict, net_settings: dict) -> dict | udpmasks = stream.get("udpmasks") if isinstance(udpmasks, list) and udpmasks and not finalmask.get("udp"): - finalmask["udp"] = deepcopy(udpmasks) + finalmask["udp"] = list(udpmasks) return finalmask or None diff --git a/app/db/__init__.py b/app/db/__init__.py index 5cc091f93..6d3b6c872 100644 --- a/app/db/__init__.py +++ b/app/db/__init__.py @@ -2,7 +2,6 @@ from .base import Base, GetDB, get_db # noqa - from .models import JWT, System, User # noqa __all__ = [ diff --git a/app/db/base.py b/app/db/base.py index 4936ce191..a985280af 100644 --- a/app/db/base.py +++ b/app/db/base.py @@ -16,8 +16,9 @@ pool_size=database_settings.pool_size, max_overflow=database_settings.max_overflow, pool_recycle=database_settings.pool_recycle, - pool_timeout=5, + pool_timeout=database_settings.pool_timeout, pool_pre_ping=True, + pool_use_lifo=True, echo=database_settings.echo_queries, ) diff --git a/app/db/migrations/versions/a1b2c3d4e5f6_add_performance_indexes.py b/app/db/migrations/versions/a1b2c3d4e5f6_add_performance_indexes.py new file mode 100644 index 000000000..8fe05af70 --- /dev/null +++ b/app/db/migrations/versions/a1b2c3d4e5f6_add_performance_indexes.py @@ -0,0 +1,48 @@ +"""add performance indexes + +Revision ID: a1b2c3d4e5f6 +Revises: fbfc49f01004 +Create Date: 2026-06-14 12:00:00.000000 + +""" + +from alembic import op + + +# revision identifiers, used by Alembic. +revision = "a1b2c3d4e5f6" +down_revision = "fbfc49f01004" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + _create_index("idx_users_status", "users", ["status"]) + _create_index("idx_users_status_expire", "users", ["status", "expire"]) + _create_index("idx_users_status_used_traffic", "users", ["status", "used_traffic"]) + _create_index("idx_nodes_status", "nodes", ["status"]) + _create_index("idx_notification_reminders_user_id", "notification_reminders", ["user_id"]) + _create_index("idx_hosts_inbound_tag", "hosts", ["inbound_tag"]) + _create_index("idx_hosts_is_disabled", "hosts", ["is_disabled"]) + _create_index("idx_temp_keys_action", "temp_keys", ["action"]) + _create_index("idx_node_stats_node_id_created_at", "node_stats", ["node_id", "created_at"]) + + +def downgrade() -> None: + _drop_index("idx_node_stats_node_id_created_at", "node_stats") + _drop_index("idx_temp_keys_action", "temp_keys") + _drop_index("idx_hosts_is_disabled", "hosts") + _drop_index("idx_hosts_inbound_tag", "hosts") + _drop_index("idx_notification_reminders_user_id", "notification_reminders") + _drop_index("idx_nodes_status", "nodes") + _drop_index("idx_users_status_used_traffic", "users") + _drop_index("idx_users_status_expire", "users") + _drop_index("idx_users_status", "users") + + +def _create_index(name: str, table_name: str, columns: list) -> None: + op.create_index(name, table_name, columns, unique=False, if_not_exists=True) + + +def _drop_index(name: str, table_name: str) -> None: + op.drop_index(name, table_name=table_name) diff --git a/app/jobs/record_usages.py b/app/jobs/record_usages.py index 76b8cd284..a09b5eb6d 100644 --- a/app/jobs/record_usages.py +++ b/app/jobs/record_usages.py @@ -98,10 +98,17 @@ def _chunked(items: list, size: int): yield items[index : index + size] +_dialect_cache: str | None = None + + async def get_dialect() -> str: - """Get the database dialect name without holding the session open.""" + """Get the database dialect name, cached after first lookup.""" + global _dialect_cache + if _dialect_cache is not None: + return _dialect_cache async with GetDB() as db: - return db.bind.dialect.name + _dialect_cache = db.bind.dialect.name + return _dialect_cache def build_node_user_usage_upsert(dialect: str, upsert_params: list[dict]): @@ -436,11 +443,8 @@ async def record_user_stats_batched(all_node_params: dict, usage_coefficients: d async def record_node_stats_batched(all_node_params: dict): """ - Record node-level statistics for ALL nodes in batched operations. - This reduces write amplification and lock contention. - - Args: - all_node_params: Dict mapping node_id -> list of node stat params + Record node-level statistics for ALL nodes in a single batched transaction. + This reduces write amplification, lock contention, and connection overhead. """ if not all_node_params: return @@ -448,35 +452,49 @@ async def record_node_stats_batched(all_node_params: dict): created_at = _get_time_bucket() dialect = await get_dialect() - # Process each node's stats with concurrency control - async def _record_single_node(node_id: int, params: list[dict]): + # Aggregate all node stats first + upsert_params = [] + for node_id, params in all_node_params.items(): if not params: - return - - # Aggregate uplink and downlink from params + continue total_up = sum(p.get("up", 0) for p in params) total_down = sum(p.get("down", 0) for p in params) - if not (total_up or total_down): - return - - upsert_param = { - "node_id": node_id, - "created_at": created_at, - "up": total_up, - "down": total_down, - } + continue + upsert_params.append( + { + "node_id": node_id, + "created_at": created_at, + "up": total_up, + "down": total_down, + } + ) - # Execute with concurrency control - async with JOB_SEM: - queries = build_node_usage_upsert(dialect, upsert_param) - for stmt, stmt_params in queries: - await safe_execute(stmt, stmt_params) + if not upsert_params: + return - # Execute all node stats with limited concurrency - tasks = [_record_single_node(node_id, params) for node_id, params in all_node_params.items()] - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) + # Execute all node stats in a single transaction + async with JOB_SEM: + try: + async with engine.begin() as conn: + for upsert_param in upsert_params: + queries = build_node_usage_upsert(dialect, upsert_param) + for stmt, stmt_params in queries: + await conn.execute(stmt, stmt_params) + except (OperationalError, DatabaseError) as err: + logger.warning("Batch failed, falling back to per-row safe_execute: %s", err) + for upsert_param in upsert_params: + queries = build_node_usage_upsert(dialect, upsert_param) + try: + for stmt, stmt_params in queries: + await safe_execute(stmt, stmt_params) + except (OperationalError, DatabaseError) as row_err: + logger.error( + "Failed to record node usage for node %s during fallback: %s", + upsert_param.get("node_id", "?"), + row_err, + exc_info=True, + ) def _process_users_stats_response(stats_response): diff --git a/app/node/__init__.py b/app/node/__init__.py index 99d316f07..c5fc310a4 100644 --- a/app/node/__init__.py +++ b/app/node/__init__.py @@ -71,26 +71,36 @@ async def get_nodes(self) -> dict[int, PasarGuardNode]: async with self._lock.reader_lock: return self._nodes - async def get_healthy_nodes(self) -> list[tuple[int, PasarGuardNode]]: + async def _get_nodes_by_health(self, expected: Health) -> list[tuple[int, PasarGuardNode]]: async with self._lock.reader_lock: - nodes: list[tuple[int, PasarGuardNode]] = [ - (id, node) for id, node in self._nodes.items() if (await node.get_health() == Health.HEALTHY) - ] - return nodes + items = list(self._nodes.items()) + + async def get_health_with_timeout(node: PasarGuardNode) -> Health: + return await asyncio.wait_for(node.get_health(), timeout=10) + + health_results = await asyncio.gather( + *(get_health_with_timeout(node) for _, node in items), + return_exceptions=True, + ) - async def get_broken_nodes(self) -> list[tuple[int, PasarGuardNode]]: async with self._lock.reader_lock: - nodes: list[tuple[int, PasarGuardNode]] = [ - (id, node) for id, node in self._nodes.items() if (await node.get_health() == Health.BROKEN) - ] - return nodes + matched = [] + for (node_id, node), health in zip(items, health_results): + if isinstance(health, Exception): + self.logger.warning("Failed to get health for node %s: %s", node_id, health) + continue + if health == expected and self._nodes.get(node_id) is node: + matched.append((node_id, node)) + return matched + + async def get_healthy_nodes(self) -> list[tuple[int, PasarGuardNode]]: + return await self._get_nodes_by_health(Health.HEALTHY) + + async def get_broken_nodes(self) -> list[tuple[int, PasarGuardNode]]: + return await self._get_nodes_by_health(Health.BROKEN) async def get_not_connected_nodes(self) -> list[tuple[int, PasarGuardNode]]: - async with self._lock.reader_lock: - nodes: list[tuple[int, PasarGuardNode]] = [ - (id, node) for id, node in self._nodes.items() if (await node.get_health() == Health.NOT_CONNECTED) - ] - return nodes + return await self._get_nodes_by_health(Health.NOT_CONNECTED) async def _snapshot_nodes(self) -> list[PasarGuardNode]: async with self._lock.reader_lock: diff --git a/config.py b/config.py index 6f7521a30..141302aee 100644 --- a/config.py +++ b/config.py @@ -32,7 +32,8 @@ class DatabaseSettings(EnvSettings): url: str = Field(default="sqlite+aiosqlite:///db.sqlite3", validation_alias="SQLALCHEMY_DATABASE_URL") pool_size: int = Field(default=25, validation_alias="SQLALCHEMY_POOL_SIZE") max_overflow: int = Field(default=60, validation_alias="SQLALCHEMY_MAX_OVERFLOW") - pool_recycle: int = Field(default=300, validation_alias="SQLALCHEMY_POOL_RECYCLE") + pool_recycle: int = Field(default=1800, ge=1, validation_alias="SQLALCHEMY_POOL_RECYCLE") + pool_timeout: int = Field(default=15, ge=1, validation_alias="SQLALCHEMY_POOL_TIMEOUT") echo_queries: bool = Field(default=False, validation_alias="ECHO_SQL_QUERIES") @cached_property