refactor(nodes): use less sessions to avoid overflow

Author: M03ED

app/db/crud/node.py

@@ -1,7 +1,7 @@
 from datetime import datetime, timezone
 from typing import Optional, Union
 
-from sqlalchemy import and_, delete, func, select, update
+from sqlalchemy import and_, delete, func, select, update, bindparam
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.db.models import Node, NodeStat, NodeStatus, NodeUsage, NodeUserUsage
@@ -278,6 +278,49 @@ def _table_model(table: UsageTable):
     raise ValueError("Invalid table enum")
 
 
+async def bulk_update_node_status(
+    db: AsyncSession,
+    updates: list[dict],
+) -> None:
+    """
+    Update multiple node statuses in a single query using bindparam.
+
+    Args:
+        db (AsyncSession): The database session.
+        updates (list[dict]): List of updates with keys: node_id, status, message, xray_version, node_version.
+
+    Example:
+        updates = [
+            {"node_id": 1, "status": NodeStatus.connected, "message": "", "xray_version": "1.8.0", "node_version": "0.1.0"},
+            {"node_id": 2, "status": NodeStatus.error, "message": "Connection failed", "xray_version": "", "node_version": ""},
+        ]
+    """
+    if not updates:
+        return
+
+    stmt = (
+        update(Node)
+        .where(Node.id == bindparam("node_id"))
+        .values(
+            status=bindparam("status"),
+            message=bindparam("message"),
+            xray_version=bindparam("xray_version"),
+            node_version=bindparam("node_version"),
+            last_status_change=bindparam("now"),
+        )
+    )
+
+    # Add timestamp to each update
+    now = datetime.now(timezone.utc)
+    for upd in updates:
+        upd["now"] = now
+
+    # Execute using connection-level execute (bypasses ORM, allows bindparam with WHERE)
+    conn = await db.connection()
+    await conn.execute(stmt, updates)
+    await db.commit()
+
+
 async def clear_usage_data(
     db: AsyncSession, table: UsageTable, start: datetime | None = None, end: datetime | None = None
 ):

app/jobs/node_checker.py

@@ -55,14 +55,21 @@ async def update_node_connection_status(node_id: int, node: PasarGuardNode):
     """
     try:
         await node.get_backend_stats(timeout=8)
-        await node_operator.update_node_status(
-            node_id, NodeStatus.connected, await node.core_version(), await node.node_version()
-        )
+        async with GetDB() as db:
+            await NodeOperation._update_single_node_status(
+                db,
+                node_id,
+                NodeStatus.connected,
+                xray_version=await node.core_version(),
+                node_version=await node.node_version(),
+            )
     except NodeAPIError as e:
         if e.code > -3:
-            await node_operator.update_node_status(node_id, NodeStatus.error, err=e.detail)
+            async with GetDB() as db:
+                await NodeOperation._update_single_node_status(db, node_id, NodeStatus.error, message=e.detail)
         if e.code > 0:
-            await node_operator.connect_node(node_id=node_id)
+            async with GetDB() as db:
+                await node_operator.connect_node_wrapper(db, node_id)
 
 
 async def process_node_health_check(db_node: Node, node: PasarGuardNode):
@@ -77,16 +84,23 @@ async def process_node_health_check(db_node: Node, node: PasarGuardNode):
         return
 
     if node.requires_hard_reset():
-        await node_operator.connect_node(db_node.id)
+        async with GetDB() as db:
+            await node_operator.connect_node_wrapper(db, db_node.id)
         return
 
     try:
         health = await asyncio.wait_for(verify_node_backend_health(node, db_node.name), timeout=15)
     except asyncio.TimeoutError:
-        await node_operator.update_node_status(db_node.id, NodeStatus.error, err="Health check timeout")
+        async with GetDB() as db:
+            await NodeOperation._update_single_node_status(
+                db, db_node.id, NodeStatus.error, message="Health check timeout"
+            )
         return
     except NodeAPIError:
-        await node_operator.update_node_status(db_node.id, NodeStatus.error, err="Get health failed")
+        async with GetDB() as db:
+            await NodeOperation._update_single_node_status(
+                db, db_node.id, NodeStatus.error, message="Get health failed"
+            )
         return
 
     # Skip nodes that are already healthy and connected
@@ -95,12 +109,14 @@ async def process_node_health_check(db_node: Node, node: PasarGuardNode):
 
     # Update status for recovering nodes
     if db_node.status in (NodeStatus.connecting, NodeStatus.error) and health == Health.HEALTHY:
-        await node_operator.update_node_status(
-            db_node.id,
-            NodeStatus.connected,
-            core_version=await node.core_version(),
-            node_version=await node.node_version(),
-        )
+        async with GetDB() as db:
+            await NodeOperation._update_single_node_status(
+                db,
+                db_node.id,
+                NodeStatus.connected,
+                xray_version=await node.core_version(),
+                node_version=await node.node_version(),
+            )
         return
 
     # For all other cases, update connection status
@@ -126,27 +142,10 @@ async def initialize_nodes():
     async with GetDB() as db:
         db_nodes = await get_nodes(db=db, enabled=True)
 
-        # Semaphore to limit concurrent node startups to 3 at a time
-        semaphore = asyncio.Semaphore(3)
-
-        async def start_node(node: Node):
-            try:
-                await node_manager.update_node(node)
-            except NodeAPIError as e:
-                await node_operator.update_node_status(node.id, NodeStatus.error, err=e.detail)
-                return
-
-            await node_operator.connect_node(node_id=node.id)
-
-        async def start_node_with_limit(node: Node):
-            async with semaphore:
-                await start_node(node)
-
         if not db_nodes:
             logger.warning("Attention: You have no node, you need to have at least one node")
         else:
-            start_tasks = [start_node_with_limit(node=db_node) for db_node in db_nodes]
-            await asyncio.gather(*start_tasks)
+            await node_operator.connect_nodes_bulk(db, db_nodes)
             logger.info("All nodes' cores have been started.")
 
     scheduler.add_job(

app/models/node.py

@@ -161,3 +161,15 @@ class NodeResponse(Node):
     message: str | None
 
     model_config = ConfigDict(from_attributes=True)
+
+
+class NodeNotification(BaseModel):
+    """Lightweight node model for sending notifications without database fetch."""
+
+    id: int
+    name: str
+    xray_version: str | None = None
+    node_version: str | None = None
+    message: str | None = None
+
+    model_config = ConfigDict(from_attributes=True)

app/notification/__init__.py

@@ -5,7 +5,7 @@
 from . import webhook as wh
 from app.models.host import BaseHost
 from app.models.user_template import UserTemplateResponse
-from app.models.node import NodeResponse
+from app.models.node import NodeNotification, NodeResponse
 from app.models.group import GroupResponse
 from app.models.core import CoreResponse
 from app.models.admin import AdminDetails
@@ -63,12 +63,12 @@ async def remove_node(node: NodeResponse, by: str):
         await asyncio.gather(ds.remove_node(node, by), tg.remove_node(node, by))
 
 
-async def connect_node(node: NodeResponse):
+async def connect_node(node: NodeNotification):
     if (await notification_enable()).node.connect:
         await asyncio.gather(ds.connect_node(node), tg.connect_node(node))
 
 
-async def error_node(node: NodeResponse):
+async def error_node(node: NodeNotification):
     if (await notification_enable()).node.error:
         await asyncio.gather(ds.error_node(node), tg.error_node(node))
 

app/notification/discord/node.py

@@ -1,7 +1,7 @@
 import copy
 
 from app.notification.client import send_discord_webhook
-from app.models.node import NodeResponse
+from app.models.node import NodeNotification, NodeResponse
 from app.models.settings import NotificationSettings
 from app.settings import notification_settings
 from app.utils.helpers import escape_ds_markdown_list, escape_ds_markdown
@@ -54,7 +54,7 @@ async def remove_node(node: NodeResponse, by: str):
         await send_discord_webhook(data, settings.discord_webhook_url)
 
 
-async def connect_node(node: NodeResponse):
+async def connect_node(node: NodeNotification):
     name = escape_ds_markdown(node.name)
     message = copy.deepcopy(messages.CONNECT_NODE)
     message["description"] = message["description"].format(
@@ -71,7 +71,7 @@ async def connect_node(node: NodeResponse):
         await send_discord_webhook(data, settings.discord_webhook_url)
 
 
-async def error_node(node: NodeResponse):
+async def error_node(node: NodeNotification):
     name, node_message = escape_ds_markdown_list((node.name, node.message))
     message = copy.deepcopy(messages.ERROR_NODE)
     message["description"] = message["description"].format(name=name, error=node_message)

app/notification/telegram/node.py

@@ -1,7 +1,7 @@
 from html import escape
 
 from app.notification.client import send_telegram_message
-from app.models.node import NodeResponse
+from app.models.node import NodeNotification, NodeResponse
 from app.models.settings import NotificationSettings
 from app.settings import notification_settings
 from app.utils.helpers import escape_tg_html
@@ -38,7 +38,7 @@ async def remove_node(node: NodeResponse, by: str):
         )
 
 
-async def connect_node(node: NodeResponse):
+async def connect_node(node: NodeNotification):
     data = messages.CONNECT_NODE.format(
         name=escape(node.name), node_version=node.node_version, core_version=node.xray_version, id=node.id
     )
@@ -49,7 +49,7 @@ async def connect_node(node: NodeResponse):
         )
 
 
-async def error_node(node: NodeResponse):
+async def error_node(node: NodeNotification):
     name, message = escape_tg_html((node.name, node.message))
     data = messages.ERROR_NODE.format(name=name, error=message, id=node.id)
     settings: NotificationSettings = await notification_settings()