coryHawkvelt
0237024d1c
Introduce `ENABLE_TASK_ASSIGNMENT_DEBUG` configuration to control verbose logging in task assignment, Redis subscription, and worker management flows. This reduces log noise by default while allowing detailed debugging when enabled.
66 lines
2.6 KiB
Python
66 lines
2.6 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
|
|
|
|
from websocket_server.config import ENABLE_TASK_ASSIGNMENT_DEBUG
|
|
|
|
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", "KEA")
|
|
logger.debug(f"Total available in pool: {len(redis_client.connection_pool._available_connections)}")
|
|
|
|
pubsub = redis_client.pubsub()
|
|
channels = [
|
|
f"__keyspace@2__:worker_status_{worker_id}", #TODO - The number 2 here is important, it's the Redis DB number which must align with the resid URL in config.py!
|
|
f"__keyspace@2__: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}")
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
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":
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Queue is true, setting dispatch flag")
|
|
worker_dispatch_flags[worker_id] = True
|
|
else:
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Queue false, no dispatch needed")
|
|
else:
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(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()
|