49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
# websocket_server/events/ack.py
|
|
|
|
from flask_socketio import SocketIO
|
|
import logging
|
|
from datetime import datetime
|
|
import json
|
|
|
|
from websocket_server.events import base
|
|
from websocket_server.config import get_db_session, get_redis_client
|
|
from websocket_server.models import Task
|
|
|
|
logger = logging.getLogger("websocket_server")
|
|
socketio = base.socketio
|
|
|
|
def register_socketio_handlers(socketio):
|
|
|
|
@socketio.on("ack")
|
|
def handle_ack(data):
|
|
"""
|
|
Worker acknowledges task completion.
|
|
Updates database with result and execution timing.
|
|
"""
|
|
worker_id = data.get("worker_id")
|
|
task_id = data.get("task_id")
|
|
result = data.get("result", {})
|
|
|
|
logger.debug(f"[{worker_id}] Acknowledged task {task_id}, result: {result}")
|
|
|
|
try:
|
|
with get_db_session() as session:
|
|
task = session.query(Task).filter(Task.id == task_id).first()
|
|
if not task:
|
|
logger.warning(f"[{worker_id}] Task {task_id} not found in DB")
|
|
return
|
|
|
|
task.status = "acknowledged"
|
|
task.success = 1 if result.get("success") else 0
|
|
task.response = json.dumps(result.get("response", ""))
|
|
task.finish_time = datetime.utcnow()
|
|
task.execution_time = (task.finish_time - task.start_time).total_seconds()
|
|
|
|
logger.info(f"[{worker_id}] Task {task_id} acknowledged. Exec time: {task.execution_time:.2f}s")
|
|
|
|
redis_client = get_redis_client()
|
|
redis_client.set(f"worker_status_{worker_id}", "idle")
|
|
|
|
except Exception as e:
|
|
logger.exception(f"[{worker_id}] Failed to acknowledge task {task_id}: {e}")
|