refactor: migrate client heartbeat management to atomic MongoDB operations and a dedicated background loop

central_api_prototype.py

@@ -176,18 +176,31 @@ def append_to_logs(log_type: str, client_id: str, detail: Any):
     except Exception as e:
         print(f"[!] Error writing log to MongoDB: {e}")
 
-def mark_client_active(client_id: str, clients: Dict[str, Any]):
+def mark_client_active(client_id: str):
-    if client_id in clients:
+    try:
-        was_active = clients[client_id].get("active", False)
+        now_str = datetime.now().isoformat()
-        if not was_active:
+        # Atomically check if transitioned from offline to online (active is not True)
-            if clients[client_id].get("last_down_alert_time"):
+        old_doc = clients_collection.find_one_and_update(
+            {"_id": client_id, "active": {"$ne": True}},
+            {"$set": {"active": True, "last_seen": now_str, "last_down_alert_time": None}},
+            return_document=False
+        )
+        if old_doc:
+            # Client transitioned from inactive/None to active!
+            # Send UP recovery alert only if they had a down alert time logged
+            if old_doc.get("last_down_alert_time"):
                 send_rocketchat_notification(
                     text=f"✅ **System UP:** Client `{client_id}` has recovered and is back online.",
                     color="#2ecc71"
                 )
-                clients[client_id]["last_down_alert_time"] = None
+        else:
-        clients[client_id]["active"] = True
+            # Client is already active, just update their last_seen timestamp
-        clients[client_id]["last_seen"] = datetime.now().isoformat()
+            clients_collection.update_one(
+                {"_id": client_id},
+                {"$set": {"last_seen": now_str}}
+            )
+    except Exception as e:
+        print(f"[!] Error marking client active: {e}")
 
 def load_clients() -> Dict[str, Any]:
     try:
@@ -195,92 +208,93 @@ def load_clients() -> Dict[str, Any]:
         clients_collection.delete_many({"_id": {"$in": ["", None]}})
         clients_collection.delete_many({"_id": {"$regex": "^\\s*$"}})
 
-
-
         # Load all documents
         cursor = clients_collection.find()
         data = {}
         for doc in cursor:
             cid = doc["_id"]
-            data[cid] = {k: v for k, v in doc.items() if k != "_id"}
-            
-        # Compute active status dynamically based on 30-second heartbeat check-ins
-        modified = False
-        for cid, info in data.items():
-            last_seen_str = info.get("last_seen")
-            was_active = info.get("active", False)
-            is_active = False
             
+            # Compute active status dynamically for reads, but do NOT write back to database or alert
+            last_seen_str = doc.get("last_seen")
+            is_active = doc.get("active", False)
             if last_seen_str:
                 try:
                     last_seen_dt = datetime.fromisoformat(last_seen_str)
-                    if (datetime.now() - last_seen_dt).total_seconds() < 30:
+                    if (datetime.now() - last_seen_dt).total_seconds() >= 30:
-                        is_active = True
+                        is_active = False
                 except Exception:
                     pass
             
-            if was_active != is_active:
+            info = {k: v for k, v in doc.items() if k != "_id"}
-                info["active"] = is_active
+            info["active"] = is_active
-                modified = True
+            data[cid] = info
-
-                # Active status transitions alerts
-                if is_active:
-                    # Transitions from offline (DOWN) to online (UP)
-                    if info.get("last_down_alert_time"):
-                        send_rocketchat_notification(
-                            text=f"✅ **System UP:** Client `{cid}` has recovered and is back online.",
-                            color="#2ecc71"
-                        )
-                        info["last_down_alert_time"] = None
-                else:
-                    # Transitions from online (UP) to offline (DOWN) for the first time
-                    now = datetime.now()
-                    info["last_down_alert_time"] = now.isoformat()
-                    send_rocketchat_notification(
-                        text=f"🚨 **System DOWN:** Client `{cid}` has missed heartbeats for over 30 seconds.",
-                        color="#e74c3c"
-                    )
-            else:
-                # No transition (state remains the same)
-                if not is_active:
-                    # Client remains offline. Check if we should send a repeating reminder.
-                    # We ONLY send a reminder if it is already tracked in last_down_alert_time.
-                    # This prevents sending alerts for historically offline systems at server boot!
-                    last_down_str = info.get("last_down_alert_time")
-                    if last_down_str:
-                        now = datetime.now()
-                        try:
-                            last_down_dt = datetime.fromisoformat(last_down_str)
-                            elapsed_seconds = (now - last_down_dt).total_seconds()
-                            if elapsed_seconds >= 3600: # 1 hour (3600 seconds)
-                                info["last_down_alert_time"] = now.isoformat()
-                                send_rocketchat_notification(
-                                    text=f"🚨 **System STILL DOWN:** Client `{cid}` remains offline (reminder sent every hour).",
-                                    color="#e74c3c"
-                                )
-                                modified = True
-                        except Exception:
-                            pass
-                
-        if modified:
-            save_clients(data)
             
         return data
     except Exception as e:
         print(f"[!] Error loading clients from MongoDB: {e}")
         return {}
 
-def save_clients(data: Dict[str, Any]):
+import asyncio
-    try:
+
-        for cid, info in data.items():
+async def monitor_heartbeats_loop():
-            # Use upsert to update existing client or insert new client automatically
+    while True:
-            clients_collection.update_one(
+        try:
-                {"_id": cid},
+            now = datetime.now()
-                {"$set": info},
+            timeout_time = (now - timedelta(seconds=30)).isoformat()
-                upsert=True
+            
-            )
+            # Find clients that are active in DB but haven't checked in for 30s
-    except Exception as e:
+            cursor = clients_collection.find({
-        print(f"[!] Error saving clients to MongoDB: {e}")
+                "active": True,
+                "last_seen": {"$lt": timeout_time}
+            })
+            for doc in cursor:
+                cid = doc["_id"]
+                now_str = now.isoformat()
+                
+                # Atomically update to active = False
+                # Ensures only one request/worker triggers the state transition and sends the alert
+                old_doc = clients_collection.find_one_and_update(
+                    {"_id": cid, "active": True, "last_seen": doc["last_seen"]},
+                    {"$set": {"active": False, "last_down_alert_time": now_str}},
+                    return_document=False
+                )
+                if old_doc:
+                    send_rocketchat_notification(
+                        text=f"🚨 **System DOWN:** Client `{cid}` has missed heartbeats for over 30 seconds.",
+                        color="#e74c3c"
+                    )
+            
+            # Hourly reminders for STILL DOWN clients
+            reminder_time = (now - timedelta(hours=1)).isoformat()
+            cursor_still_down = clients_collection.find({
+                "active": False,
+                "last_down_alert_time": {"$lt": reminder_time}
+            })
+            for doc in cursor_still_down:
+                cid = doc["_id"]
+                old_last_down = doc["last_down_alert_time"]
+                now_str = now.isoformat()
+                
+                # Atomically update reminder timestamp
+                old_doc = clients_collection.find_one_and_update(
+                    {"_id": cid, "active": False, "last_down_alert_time": old_last_down},
+                    {"$set": {"last_down_alert_time": now_str}},
+                    return_document=False
+                )
+                if old_doc:
+                    send_rocketchat_notification(
+                        text=f"🚨 **System STILL DOWN:** Client `{cid}` remains offline (reminder sent every hour).",
+                        color="#e74c3c"
+                    )
+        except Exception as e:
+            print(f"[!] Error in heartbeat monitoring loop: {e}")
+            
+        await asyncio.sleep(10)
+
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(monitor_heartbeats_loop())
+    print("[*] Heartbeat monitoring background task initialized.")
 
 class TelemetryPayload(BaseModel):
     cpu_percent: float
@@ -320,7 +334,6 @@ async def get_whitelisted_commands(platform: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed to load commands: {e}")
 
-
 @app.get("/api/get-command")
 async def get_command(client_id: str):
     """
@@ -330,25 +343,29 @@ async def get_command(client_id: str):
     if not client_id or not client_id.strip():
         raise HTTPException(status_code=400, detail="client_id cannot be empty")
         
-    clients = load_clients()
+    now_str = datetime.now().isoformat()
+    # Dynamic fleet auto-registration
+    clients_collection.update_one(
+        {"_id": client_id},
+        {"$setOnInsert": {"pending_command": "none", "active": True, "last_seen": now_str}},
+        upsert=True
+    )
     
-    # 1. Auto-register new client if we haven't seen it yet
+    # Get and reset command atomically
-    if client_id not in clients:
+    updated_doc = clients_collection.find_one_and_update(
-        print(f"[+] Dynamic fleet auto-registration: Registered new agent '{client_id}'")
+        {"_id": client_id},
-        clients[client_id] = {"pending_command": "none", "last_seen": None, "active": True}
+        {"$set": {"pending_command": "none"}},
+        return_document=False # Returns the state before update
+    )
     
-    # 2. Fetch the command
+    cmd = "none"
-    cmd = clients[client_id].get("pending_command", "none")
+    if updated_doc:
+        cmd = updated_doc.get("pending_command", "none")
         
-    # Log if a valid pending command was polled
     if cmd != "none":
         append_to_logs("command_polled", client_id, {"command": cmd})
         
-    # 3. Reset the command to "none", mark active, and update last_seen timestamp
+    mark_client_active(client_id)
-    clients[client_id]["pending_command"] = "none"
-    mark_client_active(client_id, clients)
-    
-    save_clients(clients)
     return {"command": cmd}
 
 @app.post("/api/schedule-command")
@@ -357,20 +374,18 @@ async def schedule_command(client_id: str, command: str, current_user: str = Dep
     Endpoint for your Dashboard/UI to schedule a new command.
     Supports a comma-separated list of client IDs for batch fleet updates.
     """
-    clients = load_clients()
     target_ids = [cid.strip() for cid in client_id.split(",") if cid.strip()]
-    
     if not target_ids:
         raise HTTPException(status_code=400, detail="No target client IDs specified")
 
     for cid in target_ids:
-        if cid not in clients:
+        clients_collection.update_one(
-            clients[cid] = {"pending_command": "none", "last_seen": None, "active": False}
+            {"_id": cid},
-        clients[cid]["pending_command"] = command
+            {"$set": {"pending_command": command}},
-        # Log command scheduling per system
+            upsert=True
+        )
         append_to_logs("command_scheduled", cid, {"command": command})
         
-    save_clients(clients)
     return {"status": "success", "message": f"Command '{command}' scheduled for {', '.join(target_ids)}"}
 
 @app.post("/api/telemetry")
@@ -381,24 +396,25 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload):
     if not client_id or not client_id.strip():
         raise HTTPException(status_code=400, detail="client_id cannot be empty")
         
-    clients = load_clients()
+    now_str = datetime.now().isoformat()
+    # Auto-register client if we see it through telemetry first
+    clients_collection.update_one(
+        {"_id": client_id},
+        {"$setOnInsert": {"pending_command": "none", "active": True, "last_seen": now_str}},
+        upsert=True
+    )
     
-    # Auto-register new client if we see it through telemetry first
+    mark_client_active(client_id)
-    if client_id not in clients:
-        print(f"[+] Dynamic fleet auto-registration: Registered new agent '{client_id}' through telemetry")
-        clients[client_id] = {"pending_command": "none", "last_seen": None, "active": True}
     
-    # Update last_seen timestamp, active status, and store the current telemetry details
+    # Get current alert states from database to prevent duplicate alerts
-    mark_client_active(client_id, clients)
+    client_doc = clients_collection.find_one({"_id": client_id}) or {}
-    clients[client_id]["telemetry"] = payload.dict()
+    storage_alerts = client_doc.get("active_storage_alerts", [])
+    gpu_alert_sent = client_doc.get("gpu_alert_sent", False)
     
-    # Log telemetry history event
-    append_to_logs("telemetry", client_id, payload.dict())
-
-    # Storage limits warning checker
-    storage_alerts = clients[client_id].get("active_storage_alerts", [])
     storage_modified = False
+    gpu_modified = False
 
+    # Storage warning checker
     for disk in payload.disks:
         mount = disk.get("mount", "/")
         percent = disk.get("percent", 0.0)
@@ -420,13 +436,7 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload):
                     color="#2ecc71"
                 )
 
-    if storage_modified:
-        clients[client_id]["active_storage_alerts"] = storage_alerts
-
     # GPU Failure Monitor
-    gpu_alert_sent = clients[client_id].get("gpu_alert_sent", False)
-    gpu_modified = False
-
     if len(payload.gpus) == 0:
         if not gpu_alert_sent:
             gpu_alert_sent = True
@@ -444,11 +454,39 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload):
                 color="#2ecc71"
             )
 
+    update_fields = {
+        "telemetry": payload.dict()
+    }
+    if storage_modified:
+        update_fields["active_storage_alerts"] = storage_alerts
     if gpu_modified:
-        clients[client_id]["gpu_alert_sent"] = gpu_alert_sent
+        update_fields["gpu_alert_sent"] = gpu_alert_sent
         
-    if storage_modified or gpu_modified:
+    clients_collection.update_one(
-        save_clients(clients)
+        {"_id": client_id},
+        {"$set": update_fields}
+    )
+
+    # Log telemetry history event
+    append_to_logs("telemetry", client_id, payload.dict())
+
+    data = payload.dict()
+    client_telemetry[client_id] = data
+    
+    # Log received stats
+    print(f"\n[TELEMETRY RECEIVED] from {client_id}:")
+    print(f"  CPU:  {data['cpu_percent']}%")
+    print(f"  RAM:  {data['memory_percent']}% ({data['memory_free_gb']} GB Free / {data['memory_total_gb']} GB Total)")
+    
+    for disk in data['disks']:
+        print(f"  Drive {disk['mount']}: {disk['percent']}% Used ({disk['free_gb']} GB Free / {disk['total_gb']} GB Total)")
+        
+    for gpu in data['gpus']:
+        power_str = f"Power: {gpu.get('power_draw', 'N/A')} / {gpu.get('power_limit', 'N/A')}"
+        fan_str = f"Fan: {gpu.get('fan_speed', 'N/A')}"
+        print(f"  GPU [{gpu['name']}]: Core: {gpu['utilization']}, Temp: {gpu['temp']}, VRAM: {gpu['memory_used']}/{gpu['memory_total']}, {power_str}, {fan_str}")
+        
+    return {"status": "success"}
 
     data = payload.dict()
     client_telemetry[client_id] = data