minor fixes

Author: agam1092005

state-manager/app/controller/enqueue_states.py

@@ -7,28 +7,81 @@
 from ..models.state_status_enum import StateStatusEnum
 
 from app.singletons.logs_manager import LogsManager
+from app.config.settings import get_settings
 from pymongo import ReturnDocument
 
 logger = LogsManager().get_logger()
 
 
 async def find_state(namespace_name: str, nodes: list[str]) -> State | None:
     current_time_ms = int(time.time() * 1000)
+    settings = get_settings()
+    
+    # Use pipeline to calculate timeout_at based on state-specific or global timeout
+    pipeline = [
+        {
+            "$match": {
+                "namespace_name": namespace_name,
+                "status": StateStatusEnum.CREATED,
+                "node_name": {"$in": nodes},
+                "enqueue_after": {"$lte": current_time_ms}
+            }
+        },
+        {
+            "$addFields": {
+                "status": StateStatusEnum.QUEUED,
+                "queued_at": current_time_ms,
+                "timeout_at": {
+                    "$add": [
+                        current_time_ms,
+                        {
+                            "$multiply": [
+                                {
+                                    "$ifNull": [
+                                        "$timeout_minutes", 
+                                        settings.node_timeout_minutes
+                                    ]
+                                },
+                                60000  # Convert minutes to milliseconds
+                            ]
+                        }
+                    ]
+                }
+            }
+        }
+    ]
+    
     data = await State.get_pymongo_collection().find_one_and_update(
         {
             "namespace_name": namespace_name,
             "status": StateStatusEnum.CREATED,
-            "node_name": {
-                "$in": nodes
-            },
+            "node_name": {"$in": nodes},
             "enqueue_after": {"$lte": current_time_ms}
         },
-        {
-            "$set": {
-                "status": StateStatusEnum.QUEUED,
-                "queued_at": current_time_ms
+        [
+            {
+                "$set": {
+                    "status": StateStatusEnum.QUEUED,
+                    "queued_at": current_time_ms,
+                    "timeout_at": {
+                        "$add": [
+                            current_time_ms,
+                            {
+                                "$multiply": [
+                                    {
+                                        "$ifNull": [
+                                            "$timeout_minutes", 
+                                            settings.node_timeout_minutes
+                                        ]
+                                    },
+                                    60000  # Convert minutes to milliseconds
+                                ]
+                            }
+                        ]
+                    }
+                }
             }
-        },
+        ],
         return_document=ReturnDocument.AFTER
     )
     return State(**data) if data else None

state-manager/app/models/db/state.py

@@ -29,6 +29,7 @@ class State(BaseDatabaseModel):
     fanout_id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Fanout ID of the state")
     manual_retry_fanout_id: str = Field(default="", description="Fanout ID from a manual retry request, ensuring unique retries for unite nodes.")
     queued_at: Optional[int] = Field(None, description="Unix time in milliseconds when state was queued")
+    timeout_at: Optional[int] = Field(None, description="Unix time in milliseconds when state times out")
     timeout_minutes: Optional[int] = Field(None, gt=0, description="Timeout in minutes for this specific state, taken from node registration")
 
     @before_event([Insert, Replace, Save])
@@ -108,7 +109,7 @@ class Settings:
             IndexModel(
                 [
                     ("status", 1),
-                    ("queued_at", 1),
+                    ("timeout_at", 1),
                 ],
                 name="timeout_query_index"
             )

state-manager/app/tasks/check_node_timeout.py

@@ -1,43 +1,33 @@
 import time
-from beanie.operators import Ne
 from app.models.db.state import State
 from app.models.state_status_enum import StateStatusEnum
 from app.singletons.logs_manager import LogsManager
-from app.config.settings import get_settings
 
 logger = LogsManager().get_logger()
 
 
 async def check_node_timeout():
     try:
-        settings = get_settings()
         current_time_ms = int(time.time() * 1000)
 
         logger.info(f"Checking for timed out nodes at {current_time_ms}")
 
-        # Find all QUEUED states with queued_at set
-        queued_states = await State.find(
-            State.status == StateStatusEnum.QUEUED,
-            Ne(State.queued_at, None)
-        ).to_list()
+        # Use database query to find and update timed out states in one operation
+        result = await State.get_pymongo_collection().update_many(
+            {
+                "status": StateStatusEnum.QUEUED,
+                "timeout_at": {"$ne": None, "$lte": current_time_ms}
+            },
+            {
+                "$set": {
+                    "status": StateStatusEnum.TIMEDOUT,
+                    "error": "Node execution timed out"
+                }
+            }
+        )
 
-        states_to_timeout = []
-        
-        for state in queued_states:
-            # Use state-specific timeout if available, otherwise fall back to global
-            timeout_minutes = state.timeout_minutes if state.timeout_minutes else settings.node_timeout_minutes
-            timeout_ms = timeout_minutes * 60 * 1000
-            timeout_threshold = current_time_ms - timeout_ms
-            
-            if state.queued_at <= timeout_threshold:
-                state.status = StateStatusEnum.TIMEDOUT
-                state.error = f"Node execution timed out after {timeout_minutes} minutes"
-                states_to_timeout.append(state)
-
-        if states_to_timeout:
-            # Update all timed out states in bulk
-            await State.save_all(states_to_timeout)
-            logger.info(f"Marked {len(states_to_timeout)} states as TIMEDOUT")
+        if result.modified_count > 0:
+            logger.info(f"Marked {result.modified_count} states as TIMEDOUT")
 
     except Exception:
         logger.error("Error checking node timeout", exc_info=True)