3cloud-backend/websocket_server/redis_utils.py

# 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()