kaushik/rmm-backend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
kaushik
2026-06-16 12:40:12 +05:30
parent 9987b3a45a
commit 80e1bd8e3a
1 changed files with 159 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

268
central_api_prototype.py
View File

@@ -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]):
if client_id in clients:
was_active = clients[client_id].get("active", False)
if not was_active:
if clients[client_id].get("last_down_alert_time"):
def mark_client_active(client_id: str):
try:
now_str = datetime.now().isoformat()
# Atomically check if transitioned from offline to online (active is not True)
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
clients[client_id]["active"] = True
clients[client_id]["last_seen"] = datetime.now().isoformat()
else:
# Client is already active, just update their last_seen timestamp
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:
is_active = True
if (datetime.now() - last_seen_dt).total_seconds() >= 30:
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
modified = True
# 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)
data[cid] = info
return data
except Exception as e:
print(f"[!] Error loading clients from MongoDB: {e}")
return {}
def save_clients(data: Dict[str, Any]):
import asyncio
async def monitor_heartbeats_loop():
while True:
try:
for cid, info in data.items():
# Use upsert to update existing client or insert new client automatically
clients_collection.update_one(
{"_id": cid},
{"$set": info},
upsert=True
now = datetime.now()
timeout_time = (now - timedelta(seconds=30)).isoformat()
# Find clients that are active in DB but haven't checked in for 30s
cursor = clients_collection.find({
"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 saving clients to MongoDB: {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
if client_id not in clients:
print(f"[+] Dynamic fleet auto-registration: Registered new agent '{client_id}'")
clients[client_id] = {"pending_command": "none", "last_seen": None, "active": True}
# Get and reset command atomically
updated_doc = clients_collection.find_one_and_update(
{"_id": client_id},
{"$set": {"pending_command": "none"}},
return_document=False # Returns the state before update
)
# 2. Fetch the command
cmd = clients[client_id].get("pending_command", "none")
cmd = "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
clients[client_id]["pending_command"] = "none"
mark_client_active(client_id, clients)
save_clients(clients)
mark_client_active(client_id)
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[cid] = {"pending_command": "none", "last_seen": None, "active": False}
clients[cid]["pending_command"] = command
# Log command scheduling per system
clients_collection.update_one(
{"_id": cid},
{"$set": {"pending_command": command}},
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
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}
mark_client_active(client_id)
# Update last_seen timestamp, active status, and store the current telemetry details
mark_client_active(client_id, clients)
clients[client_id]["telemetry"] = payload.dict()
# Get current alert states from database to prevent duplicate alerts
client_doc = clients_collection.find_one({"_id": client_id}) or {}
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:
save_clients(clients)
clients_collection.update_one(
{"_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