refactor(jobs): node_health_check

Author: M03ED

app/jobs/node_checker.py

@@ -1,6 +1,6 @@
 import asyncio
 
-from PasarGuardNodeBridge import NodeAPIError, PasarGuardNode
+from PasarGuardNodeBridge import NodeAPIError, PasarGuardNode, Health
 
 from app import on_shutdown, on_startup, scheduler
 from app.db import GetDB
@@ -32,13 +32,29 @@ async def check_node(id: int, node: PasarGuardNode):
             if e.code > 0:
                 await node_operator.connect_node(node_id=id)
 
-    broken_nodes, not_connected_nodes = await node_manager.get_nodes_by_health_status()
+    async def check_health(db_node: Node, node: PasarGuardNode):
+        if node is None:
+            return
+        try:
+            health = await asyncio.wait_for(node.get_health(), timeout=10)
+        except (asyncio.TimeoutError, NodeAPIError):
+            await node_operator.update_node_status(db_node.id, NodeStatus.error, err="Get health timeout")
+
+        if db_node.status in (NodeStatus.connecting, NodeStatus.error) and health is Health.HEALTHY:
+            await node_operator.update_node_status(db_node.id, NodeStatus.connected)
+
+        elif db_node.status in (NodeStatus.connecting, NodeStatus.error) and health is Health.NOT_CONNECTED:
+            await node_operator.connect_node(node_id=db_node.id)
 
-    check_tasks = [check_node(id, node) for id, node in broken_nodes]
-    connect_tasks = [node_operator.connect_node(id) for id, _ in not_connected_nodes]
+        elif db_node.status == NodeStatus.connected and health is not Health.HEALTHY:
+            await check_node(db_node.id, node)
+
+    async with GetDB() as db:
+        db_nodes = await get_nodes(db=db, enabled=True)
+        dict_nodes = await node_manager.get_nodes()
 
-    # Use return_exceptions=True to prevent one failed node from stopping others
-    await asyncio.gather(*check_tasks + connect_tasks, return_exceptions=True)
+        check_tasks = [check_health(db_node, dict_nodes[db_node.id]) for db_node in db_nodes]
+        await asyncio.gather(*check_tasks, return_exceptions=True)
 
 
 @on_startup

app/node/__init__.py

@@ -86,32 +86,6 @@ async def get_not_connected_nodes(self) -> list[tuple[int, PasarGuardNode]]:
             ]
             return nodes
 
-    async def get_nodes_by_health_status(
-        self,
-    ) -> tuple[list[tuple[int, PasarGuardNode]], list[tuple[int, PasarGuardNode]]]:
-        """Optimized method to get both broken and not_connected nodes in single lock acquisition"""
-        async with self._lock.reader_lock:
-            if not self._nodes:
-                return [], []
-
-            health_tasks = [node.get_health() for node in self._nodes.values()]
-            health_results = await asyncio.gather(*health_tasks, return_exceptions=True)
-
-            broken_nodes = []
-            not_connected_nodes = []
-
-            for (node_id, node), health in zip(self._nodes.items(), health_results):
-                if isinstance(health, Exception):
-                    continue
-                if health == Health.BROKEN:
-                    broken_nodes.append((node_id, node))
-                elif health == Health.NOT_CONNECTED:
-                    not_connected_nodes.append((node_id, node))
-                else:
-                    continue
-
-            return broken_nodes, not_connected_nodes
-
     async def update_user(self, user: UserResponse, inbounds: list[str] = None):
         proto_user = serialize_user_for_node(user.id, user.username, user.proxy_settings.dict(), inbounds)