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

# websocket_server/events/logs.py

from flask import request
from flask_socketio import emit
import logging
import json
import requests

from websocket_server.config import get_redis_client
from websocket_server.events import base
from websocket_server.shared_state import connected_sids, connected_sids_lock, connected_workers

logger = logging.getLogger("websocket_server")
socketio = base.socketio


api_server_url = "http://127.0.0.1:5000/api"
def register_socketio_handlers(socketio):

    @socketio.on("user_request_container_logs")
    def handle_user_log_request(data):
        container_id = data.get("container_id")
        user_id = data.get("user_id")
        lines = data.get("lines", 100)
        request_id = data.get("request_id")

        logger.info(f"[{request_id}] User {user_id} requested logs for container {container_id}")

        try:
            container_response = requests.get(f"{api_server_url}/workloads/containers/{container_id}")
            if container_response.status_code != 200:
                emit("user_log_response", {
                    "success": False,
                    "logs": "Container not found",
                    "request_id": request_id
                }, to=request.sid)
                return

            container_info = container_response.json()
            worker_id = container_info["workload_host_id"]

            if worker_id not in connected_workers:
                emit("user_log_response", {
                    "success": False,
                    "logs": "Worker not connected",
                    "request_id": request_id
                }, to=request.sid)
                return

            redis_client = get_redis_client()
            request_context = {
                "user_sid": request.sid,
                "request_id": request_id,
                "container_id": container_id,
                "user_id": user_id,
                "worker_id": worker_id
            }
            redis_client.setex(f"log_request_context:{request_id}", 60, json.dumps(request_context))

            socketio.emit("container-log", {
                "container_name": container_info["id"],
                "lines": lines,
                "request_id": request_id,
                "worker_id": worker_id
            }, to=connected_workers[worker_id])

            logger.info(f"[{request_id}] Dispatched log request to worker {worker_id}")

        except Exception as e:
            logger.exception(f"[{request_id}] Error processing log request")
            emit("user_log_response", {
                "success": False,
                "logs": "Internal error",
                "request_id": request_id
            }, to=request.sid)


    @socketio.on("container_log_response")
    def handle_worker_log_response(data):
        request_id = data.get("request_id")
        redis_client = get_redis_client()
        context_data = redis_client.get(f"log_request_context:{request_id}")

        if not context_data:
            logger.warning(f"[{request_id}] No context found")
            return

        try:
            context = json.loads(context_data)
            user_sid = context.get("user_sid")
            socketio.emit("user_log_response", {
                "container_id": context["container_id"],
                "logs": data.get("logs"),
                "success": data.get("success", False),
                "request_id": request_id
            }, to=user_sid)
            redis_client.delete(f"log_request_context:{request_id}")
            logger.info(f"[{request_id}] Delivered one-shot logs to user")
        except Exception as e:
            logger.exception(f"[{request_id}] Error handling worker response")


    @socketio.on("user_start_log_stream")
    def handle_user_stream_start(data):
        container_id = data.get("container_id")
        user_id = data.get("user_id")
        request_id = data.get("request_id")

        try:
            logger.info(f"[{request_id}] Starting log stream for container {container_id}")
            container_response = requests.get(f"{api_server_url}/workloads/containers/{container_id}")
            container_response.raise_for_status()
            container_info = container_response.json()

            worker_id = container_info["workload_host_id"]
            if worker_id not in connected_workers:
                emit("user_log_response", {
                    "success": False,
                    "logs": "Worker not connected",
                    "request_id": request_id
                }, to=request.sid)
                return

            redis_client = get_redis_client()
            context_payload = {
                "user_sid": request.sid,
                "container_id": container_id,
                "user_id": user_id,
                "worker_id": worker_id
            }
            redis_client.setex(f"log_request_context:{request_id}", 600, json.dumps(context_payload))

            socketio.emit("start_container_log_stream", {
                "container_name": container_info["id"],
                "request_id": request_id,
                "worker_id": worker_id
            }, to=connected_workers[worker_id])

        except Exception as e:
            logger.exception(f"[{request_id}] Failed to start log stream")


    @socketio.on("user_stop_log_stream")
    def handle_user_stream_stop(data):
        request_id = data.get("request_id")
        redis_client = get_redis_client()
        context_data = redis_client.get(f"log_request_context:{request_id}")
        if context_data:
            context = json.loads(context_data)
            worker_id = context.get("worker_id")
        else:
            logger.warning(f"[{request_id}] No Redis context, guessing worker_id from payload")
            worker_id = data.get("worker_id")

        if worker_id in connected_workers:
            socketio.emit("stop_container_log_stream", {
                "request_id": request_id
            }, to=connected_workers[worker_id])
        redis_client.delete(f"log_request_context:{request_id}")


    @socketio.on("container-log-stream")
    def handle_streamed_logs(data):
        request_id = data.get("request_id")
        logs = data.get("logs")
        worker_id = data.get("worker_id")

        redis_client = get_redis_client()
        context_data = redis_client.get(f"log_request_context:{request_id}")
        if not context_data:
            logger.warning(f"[{request_id}] No context found for streaming")
            return

        context = json.loads(context_data)
        user_sid = context.get("user_sid")
        with base.connected_sids_lock:
            if user_sid in base.connected_sids:
                socketio.emit("user_log_stream_update", {
                    "logs": logs,
                    "request_id": request_id
                }, to=user_sid)
            else:
                logger.warning(f"[{request_id}] User SID disconnected, stopping stream")
                handle_user_stream_stop({"request_id": request_id, "worker_id": worker_id})