58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
# websocket_server/redis_utils.py
|
|
|
|
import threading
|
|
import traceback
|
|
import time
|
|
from websocket_server.config import get_redis_client
|
|
from websocket_server.task_assigner import assign_task_to_worker
|
|
from websocket_server.worker_manager import notify_worker_disconnect
|
|
import logging
|
|
|
|
logger = logging.getLogger("websocket_server")
|
|
|
|
# Shared state
|
|
worker_dispatch_flags = {}
|
|
thread_stop_flags = {}
|
|
|
|
def redis_subscribe(worker_id):
|
|
"""
|
|
Redis Pub/Sub monitor for a specific worker.
|
|
Subscribes to worker status and queue keys and sets dispatch flags accordingly.
|
|
"""
|
|
redis_client = get_redis_client()
|
|
redis_client.config_set("notify-keyspace-events", "KA")
|
|
|
|
pubsub = redis_client.pubsub()
|
|
channels = [
|
|
f"__keyspace@0__:worker_status_{worker_id}",
|
|
f"__keyspace@0__:worker_queue_{worker_id}"
|
|
]
|
|
|
|
try:
|
|
pubsub.psubscribe(channels)
|
|
logger.info(f"Subscribed to Redis keyspace events: {channels}")
|
|
|
|
while not thread_stop_flags.get(worker_id, False):
|
|
message = pubsub.get_message(timeout=1.0)
|
|
if message and message["type"] in ["pmessage", "message"]:
|
|
current_status = redis_client.get(f"worker_status_{worker_id}")
|
|
logger.debug(f"[{worker_id}] Current Redis status: {current_status}")
|
|
|
|
if current_status and current_status.lower() == "idle":
|
|
if redis_client.get(f"worker_queue_{worker_id}").lower() == "true":
|
|
logger.info(f"[{worker_id}] Queue is true, setting dispatch flag")
|
|
worker_dispatch_flags[worker_id] = True
|
|
else:
|
|
logger.info(f"[{worker_id}] Queue false, no dispatch needed")
|
|
else:
|
|
logger.info(f"[{worker_id}] Status is not idle, skipping")
|
|
|
|
logger.warning(f"[{worker_id}] Redis pub/sub loop exited, disconnecting")
|
|
notify_worker_disconnect(worker_id)
|
|
|
|
except Exception as e:
|
|
error_details = traceback.format_exc()
|
|
logger.error(f"[{worker_id}] Redis pub/sub error: {e}\n{error_details}")
|
|
finally:
|
|
pubsub.close()
|