Implimented a ping\pong monitoring system work worker liveness. Will handle the situation where a websocket server crashes.

The only thing this doesnt handle is when a client goes into offline due to a transient link failure but it comes back online. We want to avoid updating the DB every time a pong is RX'd.. Perhaps we do a seocndary corss checks with offline hosts?

websocket_server/config.py

@@ -10,16 +10,23 @@ from sqlalchemy.orm import sessionmaker, declarative_base
 import pymysql
 
 # Redis configuration
-REDIS_HOST = "172.17.0.1"
-REDIS_PORT = 6379
+REDIS_URL = "redis://172.17.0.1:6379/2"
+
+
+# Configurable intervals
+PING_INTERVAL_SECONDS = 2           # How often to send PING requests
+REDIS_PING_EXPIRY_SECONDS = 30      # Redis key expiry value
+ASSIGN_INTERVAL_SECONDS = 30
+LIVENESS_CHECK_INTERVAL_SECONDS = 5 # How often to cross check connected_workers with ping results from Redis
+PING_EXPIRY_SECONDS = 3             # How long between PING and PONG before we raise the alarm
+
 
 # Database configuration
 DATABASE_URL = "mysql://root:password@172.17.0.1:3306/theapi"
 
 # Connection pool for Redis
-redis_connection_pool = redis.ConnectionPool(
-    host=REDIS_HOST,
-    port=REDIS_PORT,
+redis_connection_pool = redis.ConnectionPool.from_url(
+    REDIS_URL,
     decode_responses=True,
     max_connections=10,
     socket_timeout=5

websocket_server/events/messages.py

@@ -1,14 +1,14 @@
 # websocket_server/events/messages.py
 
+import time
 from flask import request
 from flask_socketio import emit
 import logging
-import json
 
 from websocket_server.config import get_redis_client
-from websocket_server.events import base
 from api_client.client import IaaSClient
-from websocket_server.shared_state import connected_workers, connected_sids_lock, connected_sids, worker_lock
+from websocket_server.shared_state import connected_workers, connected_sids_lock, connected_sids, worker_lock, ping_tracker
+from websocket_server.config import REDIS_PING_EXPIRY_SECONDS
 
 
 logger = logging.getLogger("websocket_server")
@@ -115,3 +115,26 @@ def register_socketio_handlers(socketio):
         except Exception as e:
             logger.exception(f"Failed to send message to {worker_id}: {e}")
             emit("error", {"message": f"Failed to send message to {worker_id}"})
+
+    @socketio.on("worker_pong")
+    def handle_worker_pong(data):
+        worker_id = data.get("worker_id")
+        pong_id = data.get("ping_id")
+        redis = get_redis_client()
+
+        if not worker_id or not pong_id:
+            logger.warning("Malformed pong received")
+            return
+
+        ping_info = ping_tracker.get(worker_id)
+        if not ping_info or ping_info["ping_id"] != pong_id:
+            logger.warning(f"[{worker_id}] Received unmatched pong")
+            return
+
+        latency_ms = int((time.time() - ping_info["sent_time"]) * 1000)
+        logger.info(f"[{worker_id}] Pong received. Latency: {latency_ms}ms")
+
+        redis.set(f"ws_liveness:{worker_id}", int(time.time()), ex=REDIS_PING_EXPIRY_SECONDS)
+        redis.set(f"ws_latency:{worker_id}", latency_ms, ex=REDIS_PING_EXPIRY_SECONDS)
+
+        del ping_tracker[worker_id]

websocket_server/shared_state.py

@@ -6,3 +6,4 @@ connected_workers = {}
 worker_lock = threading.Lock()
 connected_sids = set()
 connected_sids_lock = threading.Lock()
+ping_tracker = {}           # Maps worker_id → {"ping_id": str, "sent_time": float}

websocket_server/worker_manager.py

@@ -1,12 +1,15 @@
 # websocket_server/worker_manager.py
 
-import threading
 import requests
+import uuid
+import time
+import threading
 import logging
-# from websocket_server.redis_utils import redis_subscribe, thread_stop_flags, worker_dispatch_flags
+from websocket_server.shared_state import connected_workers, worker_lock, ping_tracker
+from websocket_server.config import get_redis_client,PING_INTERVAL_SECONDS, ASSIGN_INTERVAL_SECONDS, LIVENESS_CHECK_INTERVAL_SECONDS, PING_EXPIRY_SECONDS
 from websocket_server.task_assigner import assign_task_to_worker
-from websocket_server.config import get_redis_client
-from websocket_server.shared_state import connected_workers, worker_lock
+from websocket_server.events import base
+
 
 logger = logging.getLogger("websocket_server")
 
@@ -16,6 +19,13 @@ api_server_url = "http://172.17.0.1:5000/api"
 # Tracks threads for each worker
 worker_threads = {}
 
+redis = get_redis_client()
+
+# Last timestamps
+last_ping_time = 0
+last_assign_time = 0
+
+
 def notify_worker_online(worker_id):
     """
     Inform the API server that a worker has connected.
@@ -65,15 +75,48 @@ def worker_dispatch_flag_check():
 
 def all_worker_watchdog():
     """
-    Periodically iterate over connected workers and ensure they are polled for task assignments.
-    Prevents idle starvation or missed Redis events.
+    Watchdog loop that:
+    - Sends a liveness ping every `PING_INTERVAL_SECONDS`
+    - Triggers task reassignment every `ASSIGN_INTERVAL_SECONDS`
     """
+    global last_ping_time, last_assign_time
     logger.info("Watchdog thread started")
+
+    redis = get_redis_client()
+    last_liveness_check = 0
+
     while True:
-        for wid in list(connected_workers.keys()):
-            logger.info(f"[{wid}] Watchdog triggering assign check")
-            assign_task_to_worker(wid)
-        threading.Event().wait(30)
+        now = time.time()
+        worker_ids = list(connected_workers.keys())
+
+        # 1. Send custom ping to all workers
+        if now - last_ping_time >= PING_INTERVAL_SECONDS:
+            if worker_ids:
+                logger.debug("Sending liveness pings to all workers")
+                for wid in worker_ids:
+                    ping_id = str(uuid.uuid4())
+                    ping_tracker[wid] = {"ping_id": ping_id, "sent_time": now}
+                    sid = connected_workers.get(wid)
+                    if sid:
+                        try:
+                            base.socketio.emit("worker_ping", {"ping_id": ping_id}, to=sid)
+                            logger.info(f"[{wid}] Sent ping ID: {ping_id} to {wid}")
+                        except Exception as e:
+                            logger.warning(f"[{wid}] Ping send failed: {e}")
+            last_ping_time = now
+
+        # 2. Trigger task assignment if due
+        if now - last_assign_time >= ASSIGN_INTERVAL_SECONDS:
+            if worker_ids:
+                logger.debug("Checking for task assignment across all workers")
+                for wid in worker_ids:
+                    logger.info(f"[{wid}] Triggering assign check")
+                    assign_task_to_worker(wid)
+            last_assign_time = now
+
+       
+
+        time.sleep(1)
 
 
 def start_worker_thread(worker_id):