3cloud-backend/websocket_server/events/messages.py

# websocket_server/events/messages.py

import time
from flask import request
from flask_socketio import emit
import logging

from websocket_server.config import get_redis_client
from api_client.client import IaaSClient
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, ENABLE_WEBSOCKET_PING_DEBUG


logger = logging.getLogger("websocket_server")

def register_socketio_handlers(socketio):
    @socketio.on("connect")
    def handle_connect():
        with connected_sids_lock:
            connected_sids.add(request.sid)
        logger.info(f"Client connected from {request.remote_addr} SID={request.sid}")
        emit("welcome", {"message": f"Connected to server. Your SID is {request.sid}"})


    @socketio.on("disconnect")
    def handle_disconnect():
        from websocket_server.worker_manager import notify_worker_disconnect

        sid = request.sid
        logger.info(f"SID={sid} disconnected")

        with connected_sids_lock:
            connected_sids.discard(sid)

        # Check if it's a worker
        worker_id = None
        with worker_lock:
            for w_id, s_id in connected_workers.items():
                if s_id == sid:
                    worker_id = w_id
                    break
            if worker_id:
                del connected_workers[worker_id]
                logger.warning(f"Worker {worker_id} disconnected")

                # Delete Redis ping keys to prevent stale heartbeats from marking worker online
                redis_client = get_redis_client()
                redis_client.delete(f"ws_liveness:{worker_id}")
                redis_client.delete(f"ws_latency:{worker_id}")

                notify_worker_disconnect(worker_id)


    @socketio.on("join_request")
    def handle_join_request(data):
        from websocket_server.worker_manager import notify_worker_online, start_worker_thread
        from websocket_server.task_assigner import assign_task_to_worker 

        worker_id = data.get("worker_id")
        worker_secret = data.get("worker_secret")

        logger.info(f"Join request received from {worker_id}")

        try:
            api = IaaSClient("xxx")
            db_worker = api.get_workload_host(worker_id)

            if not db_worker:
                logger.error(f"{worker_id} not found in DB")
                emit("join_reject", {"message": "Join request rejected"}, to=request.sid)
                return

            if db_worker["secret_key"] != worker_secret:
                logger.error(f"Invalid secret for worker {worker_id}")
                emit("join_reject", {"message": "Join request rejected"}, to=request.sid)
                return

            with worker_lock:
                connected_workers[worker_id] = request.sid

            redis_client = get_redis_client()
            redis_client.set(f"worker_status_{worker_id}", "idle")

            start_worker_thread(worker_id)
            notify_worker_online(worker_id)

            logger.info(f"Worker {worker_id} accepted")
            emit("join_accept", {"message": f"Worker {worker_id} accepted"}, to=request.sid)
            emit("join_broadcast", {"message": f"Worker {worker_id} subscribed to task queue."})
            assign_task_to_worker(worker_id)

        except Exception as e:
            logger.exception(f"Error processing join for {worker_id}: {e}")
            emit("join_reject", {"message": "Internal error during join"}, to=request.sid)


    @socketio.on("send_message")
    def handle_send_message(data):
        """
        Relay a message from a client to a specific worker by worker_id.
        """
        worker_id = data.get("worker_id")
        message = data.get("message")

        if not worker_id or not message:
            logger.warning(f"Invalid send_message request: {data}")
            emit("error", {"message": "Missing worker_id or message"})
            return

        with worker_lock:
            target_sid = connected_workers.get(worker_id)

        if not target_sid:
            logger.error(f"Worker {worker_id} is not connected")
            emit("error", {"message": f"Worker {worker_id} is not connected"})
            return

        try:
            logger.info(f"Relaying message to Worker {worker_id} (SID={target_sid})")
            emit("message", {"message": message}, to=target_sid)
        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)
        if ENABLE_WEBSOCKET_PING_DEBUG:
            logger.debug(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]