coryHawkvelt
070429ef46
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?
445 lines
17 KiB
Python
445 lines
17 KiB
Python
#Worker.py
|
|
import functools
|
|
from logger import logger
|
|
|
|
from vncsession import VNCSession
|
|
from worker_tasks.container import ContainerTask
|
|
from worker_tasks.file_presence import FilePresenceTask
|
|
from worker_tasks.ping import PingTask
|
|
from worker_tasks.ping import PingTask
|
|
from worker_tasks.report import ReportTask
|
|
from worker_tasks.ovs_sdn import OVS_SDN
|
|
from settings import settings # Import global settings
|
|
import json
|
|
import asyncio
|
|
import socketio
|
|
import threading
|
|
import time
|
|
class WorkerClient:
|
|
def __init__(self, event_queue=None, docker_monitor=None):
|
|
"""Initializes the WorkerClient using global settings."""
|
|
try:
|
|
self.worker_id = settings.get_value("WORKER_ID")
|
|
self.worker_secret = settings.get_value("WORKER_SECRET")
|
|
self.server_url = settings.get_value("WEBSOCKET_SERVER_URL")
|
|
self.debug_socketio = settings.get_value("DEBUG_SOCKETIO", False) # Get debug flag
|
|
# Store docker_monitor for container tasks
|
|
self.docker_monitor = docker_monitor
|
|
# Libvirt config is now also potentially in settings, or keep defaults here
|
|
self.libvirt_config = {
|
|
"default_volume_path": settings.get_value("DEFAULT_VOLUME_PATH", "/tmp")
|
|
}
|
|
except KeyError as e:
|
|
logger.error(f"Failed to initialize WorkerClient: Missing required setting: {e}")
|
|
raise ValueError(f"Missing required setting for WorkerClient: {e}") from e
|
|
|
|
self.event_queue = event_queue
|
|
self.running = True
|
|
|
|
|
|
if self.debug_socketio:
|
|
self.sio = socketio.AsyncClient(logger=logger, engineio_logger=logger)
|
|
else:
|
|
self.sio = socketio.AsyncClient(logger=False, engineio_logger=False)
|
|
self.joined_server = False
|
|
|
|
# Bind events
|
|
self.sio.on("connect", self.on_connect)
|
|
self.sio.on("connect_error", self.on_connect_error)
|
|
self.sio.on("disconnect", self.on_disconnect)
|
|
self.sio.on("task", self.handle_task)
|
|
self.sio.on("join_accept", self.on_join_accept)
|
|
self.sio.on("join_reject", self.on_join_reject)
|
|
self.sio.on("message", self.on_message)
|
|
# self.sio.on("*", self.debug_all_events)
|
|
|
|
|
|
self.sio.on("container-log",self.logs_retrieve)
|
|
|
|
|
|
self.log_stream_threads = {}
|
|
self.sio.on("start_container_log_stream", self.logs_stream_start)
|
|
self.sio.on("stop_container_log_stream", self.logs_stream_stop)
|
|
|
|
self.terminal_sessions = {}
|
|
self.sio.on("start_terminal", self.start_terminal_session)
|
|
self.sio.on("terminal_input", self.terminal_input)
|
|
self.sio.on("stop_terminal", self.stop_terminal_session)
|
|
|
|
|
|
|
|
self.vnc_sessions = {}
|
|
self.sio.on("start_vnc_stream_on_worker", self.start_vnc_stream)
|
|
self.sio.on("stop_vnc_stream_on_worker", self.stop_vnc_stream)
|
|
self.sio.on("vnc_frame_from_novnc", self.vnc_frame_from_novnc) # Add this line
|
|
|
|
self.sio.on("worker_ping", self.handle_worker_ping)
|
|
|
|
async def handle_worker_ping(self, data):
|
|
"""
|
|
Respond to a liveness ping from the WebSocket server.
|
|
Emits a 'worker_pong' event with the same ping_id and this worker's ID.
|
|
"""
|
|
try:
|
|
ping_id = data["ping_id"]
|
|
worker_id = self.worker_id
|
|
await self.sio.emit("worker_pong", {"ping_id": ping_id,"worker_id": worker_id})
|
|
logger.debug(f"[{worker_id}] Responded to ping {ping_id}")
|
|
except Exception as e:
|
|
logger.error(f"Error handling worker_ping: {e}")
|
|
|
|
async def on_connect(self):
|
|
logger.info("Connected to the server, requesting to join.")
|
|
await self.send_join_request()
|
|
|
|
async def on_connect_error(self, data):
|
|
logger.error(f"Connection failed: {data}")
|
|
self.joined_server = False
|
|
|
|
async def on_disconnect(self):
|
|
logger.info("Disconnected from the server.")
|
|
self.joined_server = False
|
|
|
|
async def on_message(self, data):
|
|
logger.info(f"Message received: {data}")
|
|
|
|
async def on_join_accept(self, data):
|
|
logger.info("Join accepted.")
|
|
self.joined_server = True
|
|
|
|
async def on_join_reject(self, data):
|
|
logger.error("Join rejected.")
|
|
self.joined_server = False
|
|
exit()
|
|
|
|
async def send_join_request(self):
|
|
"""Notify the server about this worker (Initial join)."""
|
|
logger.info("Sending join request")
|
|
if not self.joined_server:
|
|
await self.sio.emit("join_request", {"worker_id": self.worker_id, "worker_secret": self.worker_secret})
|
|
logger.info(f"Worker {self.worker_id} asked to join the server.")
|
|
|
|
async def handle_task(self, data):
|
|
"""Handle incoming tasks from the server."""
|
|
if not self.joined_server:
|
|
logger.error("Not joined to server. ignoring task")
|
|
return
|
|
|
|
task_type = data["type"]
|
|
task_id = data["task_id"]
|
|
job_details = json.loads(data["job_details"])
|
|
task_worker_id = data["worker_id"]
|
|
|
|
if self.docker_monitor and task_type.startswith("container-"):
|
|
job_details["docker_monitor"] = self.docker_monitor
|
|
logger.debug("Added docker_monitor to job_details for container task")
|
|
|
|
logger.info(f"Received task {task_id} of type '{task_type}' with job_details: {job_details}")
|
|
|
|
loop = asyncio.get_event_loop()
|
|
result = None
|
|
|
|
try:
|
|
# ----------------------- Task dispatch ------------------------ #
|
|
if task_type == "report":
|
|
result = await loop.run_in_executor(None, functools.partial(ReportTask("", logger).Execute))
|
|
elif task_type == "ping":
|
|
result = await loop.run_in_executor(None, functools.partial(PingTask(job_details, logger).Execute))
|
|
elif task_type == "sdn-update":
|
|
result = await loop.run_in_executor(None, functools.partial(OVS_SDN(job_details, logger).Execute))
|
|
elif task_type == "file_presence":
|
|
result = await loop.run_in_executor(None, functools.partial(FilePresenceTask(job_details, logger).Execute))
|
|
elif task_type == "pod-update":
|
|
result = await loop.run_in_executor(None, functools.partial(ContainerTask(logger).handle_pod_update, job_details))
|
|
elif task_type == "virtual-machine-create":
|
|
from worker_tasks.libvirt import LibvirtVirtualMachineTask
|
|
result = await loop.run_in_executor(None, functools.partial(LibvirtVirtualMachineTask(job_details, logger).execute))
|
|
elif task_type == "virtual-machine-delete":
|
|
from worker_tasks.libvirt import LibvirtVirtualMachineTask
|
|
result = await loop.run_in_executor(None, functools.partial(LibvirtVirtualMachineTask(job_details, logger).execute))
|
|
else:
|
|
raise ValueError(f"Unknown task type: {task_type}")
|
|
|
|
# --------------------- ACK back to server --------------------- #
|
|
await self.send_task_result(task_id, result, task_worker_id)
|
|
|
|
# ---------------- Emit launch-failure events ------------------ #
|
|
if task_type == "pod-update" and result and result.get("launch_failures"):
|
|
for fail in result["launch_failures"]:
|
|
failed_id = fail["id"]
|
|
error_info = fail["error"]
|
|
|
|
event_data=({
|
|
"source": "docker",
|
|
"event_type": "launch_failed",
|
|
"container_id": failed_id,
|
|
"system_container_id": failed_id,
|
|
"container_name": failed_id,
|
|
"timestamp": time.time(),
|
|
"error": error_info
|
|
})
|
|
|
|
# Put event in the queue for the worker client
|
|
self.event_queue.put(event_data)
|
|
logger.info(f"Added Docker event to queue: launch_failed for container {failed_id}.")
|
|
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing task {task_id}: {e}")
|
|
await self.send_task_result(task_id, {"success": False, "response": str(e)}, task_worker_id)
|
|
|
|
async def send_task_result(self, task_id, result, worker_id):
|
|
"""Send task result back to the server."""
|
|
await self.sio.emit("ack", {"task_id": task_id, "worker_id": worker_id, "result": result})
|
|
logger.debug(f"Sent result for task {task_id}: {result}")
|
|
|
|
async def send_event(self, event_data):
|
|
"""Send monitoring event to the server, handling libvirt and docker events separately."""
|
|
if not self.joined_server:
|
|
return
|
|
|
|
source = event_data.get('source', 'unknown')
|
|
event_type = event_data.get('event_type', 'event')
|
|
|
|
# Construct the base event payload
|
|
payload = {
|
|
"worker_id": self.worker_id,
|
|
"type": f"{source}_{event_type}",
|
|
"details": event_data
|
|
}
|
|
|
|
# Handle libvirt events
|
|
if source == "libvirt":
|
|
payload["details"] = {
|
|
"system_vm_id": event_data.get("system_vm_id"),
|
|
"libvirt_vm_id": event_data.get("libvirt_vm_id"),
|
|
"timestamp": event_data.get("timestamp"),
|
|
}
|
|
|
|
# Emit the event to the server
|
|
await self.sio.emit("libvirt_event", payload)
|
|
logger.debug(f"Sent libvirt event: {event_type}\n{payload}")
|
|
|
|
|
|
# Handle docker events
|
|
elif source == "docker":
|
|
payload["details"] = {
|
|
"container_id": event_data.get("container_id"),
|
|
"system_container_id": event_data.get("system_container_id"),
|
|
"container_name": event_data.get("container_name"),
|
|
"timestamp": event_data.get("timestamp"),
|
|
"status": event_data.get("details", {}).get("status"),
|
|
"image": event_data.get("details", {}).get("from"),
|
|
"action": event_data.get("details", {}).get("Action"),
|
|
**({"error": event_data.get("error")} if event_data.get("error") else {})
|
|
}
|
|
|
|
# Emit the event to the server
|
|
await self.sio.emit("docker_event", payload)
|
|
logger.debug(f"Sent {source} event: {event_type}\n{payload}")
|
|
|
|
async def process_event_queue(self):
|
|
"""Process events from the queue and send them to the server."""
|
|
if not self.event_queue:
|
|
return
|
|
|
|
while self.running:
|
|
try:
|
|
# Non-blocking check for events
|
|
if not self.event_queue.empty():
|
|
event = self.event_queue.get(block=False)
|
|
logger.debug(f"Processing event from queue: {event}")
|
|
await self.send_event(event)
|
|
await asyncio.sleep(0.1) # Small delay to prevent CPU thrashing
|
|
except Exception as e:
|
|
logger.error(f"Error in event processing: {e}")
|
|
|
|
async def start(self):
|
|
"""Worker connects to the API server and processes tasks."""
|
|
try:
|
|
logger.info(f"Worker {self.worker_id} connecting to {self.server_url}...")
|
|
await self.sio.connect(self.server_url)
|
|
|
|
# Start the event queue processor if we have an event queue
|
|
if self.event_queue:
|
|
event_processor = asyncio.create_task(self.process_event_queue())
|
|
await self.sio.wait()
|
|
event_processor.cancel()
|
|
else:
|
|
await self.sio.wait()
|
|
except Exception as e:
|
|
logger.error(f"Error: {e}")
|
|
await self.sio.disconnect()
|
|
|
|
def stop(self):
|
|
"""Stop the worker client."""
|
|
self.running = False
|
|
|
|
|
|
|
|
async def logs_retrieve(self,data):
|
|
logger.info(f"Handling container-log {data}")
|
|
|
|
container_name = data.get("container_name")
|
|
lines = int(data.get("lines", 100))
|
|
request_id = data.get("request_id")
|
|
worker_id = data.get("worker_id")
|
|
|
|
if worker_id != self.worker_id:
|
|
logger.debug(f"Ignoring container-log request for worker {worker_id}")
|
|
return
|
|
|
|
logger.info(f"Handling container-log for {container_name}, lines={lines}")
|
|
|
|
log_data=ContainerTask(logger).handle_container_log_request(data)
|
|
|
|
log_data["worker_id"] = self.worker_id
|
|
log_data["request_id"] = request_id
|
|
|
|
await self.sio.emit("container_log_response", log_data)
|
|
logger.debug(f"Sent container-log-response for {container_name}")
|
|
|
|
async def logs_stream_start(self, data):
|
|
request_id = data["request_id"]
|
|
container_name = data["container_name"]
|
|
worker_id = data.get("worker_id")
|
|
logger.debug(f"Starting stream with {data}")
|
|
|
|
if worker_id != self.worker_id:
|
|
logger.debug(f"Ignoring stream request for worker {worker_id}")
|
|
return
|
|
|
|
if request_id in self.log_stream_threads:
|
|
logger.warning(f"Log stream already running for request {request_id}")
|
|
return
|
|
|
|
# Set cancellation flag
|
|
self.log_stream_threads[request_id] = True
|
|
|
|
container_task = ContainerTask(logger)
|
|
|
|
def cancel_flag():
|
|
return not self.log_stream_threads.get(request_id, False)
|
|
|
|
def on_log(log_line):
|
|
asyncio.run(self.sio.emit("container-log-stream", {
|
|
"request_id": request_id,
|
|
"worker_id": self.worker_id,
|
|
"logs": log_line
|
|
}))
|
|
|
|
def on_error(error_msg):
|
|
asyncio.run(self.sio.emit("container-log-stream", {
|
|
"request_id": request_id,
|
|
"worker_id": self.worker_id,
|
|
"logs": f"Stream error: {error_msg}"
|
|
}))
|
|
|
|
thread = threading.Thread(
|
|
target=container_task.stream_logs,
|
|
args=(container_name, on_log, on_error, cancel_flag),
|
|
daemon=True
|
|
)
|
|
thread.start()
|
|
|
|
logger.info(f"Started log stream thread for request {request_id}")
|
|
|
|
async def logs_stream_stop(self, data):
|
|
request_id = data.get("request_id")
|
|
if request_id in self.log_stream_threads:
|
|
self.log_stream_threads[request_id] = False
|
|
logger.info(f"Stopped log stream for request {request_id}")
|
|
else:
|
|
logger.warning(f"No active log stream found for request {request_id}")
|
|
|
|
|
|
async def start_terminal_session(self, data):
|
|
request_id = data["request_id"]
|
|
container_name = data["container_name"]
|
|
command = data.get("command", "/bin/sh")
|
|
logger.debug(f"Starting terminal session with {data}")
|
|
|
|
if request_id in self.terminal_sessions:
|
|
logger.warning(f"Terminal session already exists for {request_id}")
|
|
return
|
|
|
|
container_task = ContainerTask(logger)
|
|
session = container_task.start_terminal(logger, container_name, command)
|
|
|
|
def cancel_flag():
|
|
return request_id not in self.terminal_sessions
|
|
|
|
def on_output(output):
|
|
asyncio.run(self.sio.emit("terminal_output", {
|
|
"request_id": request_id,
|
|
"data": output
|
|
}))
|
|
|
|
self.terminal_sessions[request_id] = session
|
|
|
|
threading.Thread(
|
|
target=session.stream_output,
|
|
args=(logger, on_output, cancel_flag),
|
|
daemon=True
|
|
).start()
|
|
|
|
logger.info(f"Started terminal session for {request_id}")
|
|
|
|
async def terminal_input(self, data):
|
|
logger.debug(f"Terminal input {data}")
|
|
request_id = data["request_id"]
|
|
session = self.terminal_sessions.get(request_id)
|
|
if session:
|
|
session.write(data["data"])
|
|
|
|
async def stop_terminal_session(self, data):
|
|
logger.debug(f"Stopping terminal session with {data}")
|
|
request_id = data["request_id"]
|
|
session = self.terminal_sessions.pop(request_id, None)
|
|
if session:
|
|
session.close()
|
|
logger.info(f"Closed terminal session for {request_id}")
|
|
|
|
|
|
|
|
async def start_vnc_stream(self, data):
|
|
logger.debug(f"start vnc {data}")
|
|
request_id = data["vnc_request_id"]
|
|
vnc_port = int(data["vnc_port"])
|
|
vnc_host = "127.0.0.1"
|
|
|
|
logger.info(f"[VNC] Starting session {request_id} on port {vnc_port}")
|
|
|
|
if request_id in self.vnc_sessions:
|
|
logger.warning(f"[VNC] Session already running for {request_id}")
|
|
return
|
|
|
|
def cancel_flag():
|
|
return request_id not in self.vnc_sessions
|
|
|
|
def on_output(vnc_data):
|
|
asyncio.run(self.sio.emit("vnc_frame_from_worker", {
|
|
"vnc_request_id": request_id,
|
|
"data": vnc_data
|
|
}))
|
|
|
|
session = VNCSession(logger, request_id, vnc_host, vnc_port, on_output, cancel_flag)
|
|
self.vnc_sessions[request_id] = session
|
|
|
|
thread = threading.Thread(target=session.start, daemon=True)
|
|
thread.start()
|
|
logger.info(f"[VNC] Session thread started for {request_id}")
|
|
|
|
async def vnc_frame_from_novnc(self, data):
|
|
request_id = data["vnc_request_id"]
|
|
session = self.vnc_sessions.get(request_id)
|
|
if session:
|
|
session.write(data["data"])
|
|
|
|
async def stop_vnc_stream(self, data):
|
|
request_id = data["vnc_request_id"]
|
|
session = self.vnc_sessions.pop(request_id, None)
|
|
if session:
|
|
session.close()
|
|
logger.info(f"[VNC] Stopped session for {request_id}")
|