89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
# websocket_server/events/messages.py
|
|
|
|
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
|
|
|
|
|
|
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")
|
|
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)
|