feat: add max message size calculation based on active users count

x0sina · x0sina · commit e2654ff96075 · 2026-01-05T16:40:36.000+03:30
- Introduced a new function `calculate_max_message_size` to compute the maximum message size based on the number of active users.
- Updated `NodeManager`'s `update_node` method to accept an optional `max_message_size` parameter.
- Modified `NodeOperation` to calculate and pass the max message size when updating nodes, ensuring efficient memory usage.
diff --git a/app/node/__init__.py b/app/node/__init__.py
@@ -14,6 +14,29 @@
 }
 
 
+def calculate_max_message_size(active_users_count: int) -> int:
+    """
+    Calculate max_message_size based on active users count.
+    
+    Formula: Base 10MB + (active_users_count * 50KB per user)
+    This ensures sufficient buffer for large node configurations with many users.
+    
+    Args:
+        active_users_count: Number of active users in the system
+        
+    Returns:
+        int: Max message size in bytes
+    """
+    base_size = 10 * 1024 * 1024  # 10MB base
+    per_user_size = 50 * 1024  # 50KB per user
+    calculated_size = base_size + (active_users_count * per_user_size)
+    
+    # Cap at 100MB to prevent excessive memory usage
+    max_cap = 100 * 1024 * 1024  # 100MB
+    
+    return min(calculated_size, max_cap)
+
+
 class NodeManager:
     def __init__(self):
         self._nodes: dict[int, PasarGuardNode] = {}
@@ -30,7 +53,7 @@ async def _shutdown_node(self, node: PasarGuardNode | None):
         except Exception:
             pass
 
-    async def update_node(self, node: Node) -> PasarGuardNode:
+    async def update_node(self, node: Node, max_message_size: int | None = None) -> PasarGuardNode:
         async with self._lock.writer_lock:
             old_node: PasarGuardNode | None = self._nodes.pop(node.id, None)
 
@@ -45,6 +68,7 @@ async def update_node(self, node: Node) -> PasarGuardNode:
                 logger=self.logger,
                 default_timeout=node.default_timeout,
                 internal_timeout=node.internal_timeout,
+                max_message_size=max_message_size,
                 extra={"id": node.id, "usage_coefficient": node.usage_coefficient},
             )
 
@@ -117,4 +141,4 @@ async def remove_user(self, user: UserResponse):
 node_manager: NodeManager = NodeManager()
 
 
-__all__ = ["core_users", "node_manager"]
+__all__ = ["calculate_max_message_size", "core_users", "node_manager"]
diff --git a/app/operation/node.py b/app/operation/node.py
@@ -21,8 +21,8 @@
     reset_node_usage,
     update_node_status,
 )
-from app.db.crud.user import get_user
-from app.db.models import Node, NodeStatus
+from app.db.crud.user import get_user, get_users_count_by_status
+from app.db.models import Node, NodeStatus, UserStatus
 from app.models.admin import AdminDetails
 from app.models.node import (
     NodeCoreUpdate,
@@ -37,7 +37,7 @@
     UserIPListAll,
 )
 from app.models.stats import NodeRealtimeStats, NodeStatsList, NodeUsageStatsList, Period
-from app.node import core_users, node_manager
+from app.node import calculate_max_message_size, core_users, node_manager
 from app.operation import BaseOperation
 from app.utils.logger import get_logger
 
@@ -197,8 +197,13 @@ async def create_node(self, db: AsyncSession, new_node: NodeCreate, admin: Admin
         except IntegrityError:
             await self.raise_error(message=f'Node "{new_node.name}" already exists', code=409, db=db)
 
+        # Calculate max_message_size based on active users count
+        user_counts = await get_users_count_by_status(db, [UserStatus.active])
+        active_users_count = user_counts.get(UserStatus.active.value, 0)
+        max_message_size = calculate_max_message_size(active_users_count)
+
         try:
-            await node_manager.update_node(db_node)
+            await node_manager.update_node(db_node, max_message_size=max_message_size)
             asyncio.create_task(self.connect_single_node(db, db_node.id))
         except NodeAPIError as e:
             await self._update_single_node_status(db, db_node.id, NodeStatus.error, message=e.detail)
@@ -226,8 +231,13 @@ async def modify_node(
         if db_node.status in (NodeStatus.disabled, NodeStatus.limited):
             await self.disconnect_single_node(db_node.id)
         else:
+            # Calculate max_message_size based on active users count
+            user_counts = await get_users_count_by_status(db, [UserStatus.active])
+            active_users_count = user_counts.get(UserStatus.active.value, 0)
+            max_message_size = calculate_max_message_size(active_users_count)
+
             try:
-                await node_manager.update_node(db_node)
+                await node_manager.update_node(db_node, max_message_size=max_message_size)
                 asyncio.create_task(self.connect_single_node(db, db_node.id))
             except NodeAPIError as e:
                 await self._update_single_node_status(db, db_node.id, NodeStatus.error, message=e.detail)
@@ -300,12 +310,17 @@ async def connect_nodes_bulk(
         # Fetch users ONCE for all nodes
         users = await core_users(db=db)
 
+        # Calculate max_message_size based on active users count (once for all nodes)
+        user_counts = await get_users_count_by_status(db, [UserStatus.active])
+        active_users_count = user_counts.get(UserStatus.active.value, 0)
+        max_message_size = calculate_max_message_size(active_users_count)
+
         async def connect_single(node: Node) -> dict | None:
             if node is None or node.status in (NodeStatus.disabled, NodeStatus.limited):
                 return
 
             try:
-                await node_manager.update_node(node)
+                await node_manager.update_node(node, max_message_size=max_message_size)
             except NodeAPIError as e:
                 return {
                     "node_id": node.id,
@@ -376,9 +391,14 @@ async def connect_single_node(self, db: AsyncSession, node_id: int) -> None:
         # Get core users once
         users = await core_users(db=db)
 
+        # Calculate max_message_size based on active users count
+        user_counts = await get_users_count_by_status(db, [UserStatus.active])
+        active_users_count = user_counts.get(UserStatus.active.value, 0)
+        max_message_size = calculate_max_message_size(active_users_count)
+
         # Update node manager
         try:
-            await node_manager.update_node(db_node)
+            await node_manager.update_node(db_node, max_message_size=max_message_size)
         except NodeAPIError as e:
             # Update status to error using simple CRUD
             await update_node_status(
