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.
118 lines
3.1 KiB
Python
118 lines
3.1 KiB
Python
# websocket_server/config.py
|
|
|
|
import os
|
|
import redis
|
|
import logging
|
|
from logging import StreamHandler, FileHandler
|
|
from colorlog import ColoredFormatter
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker, declarative_base
|
|
import pymysql
|
|
|
|
# Redis configuration
|
|
REDIS_URL = "redis://172.17.0.1:6379/2"
|
|
|
|
|
|
# Configurable intervals
|
|
PING_INTERVAL_SECONDS = 15 # How often to send PING requests
|
|
REDIS_PING_EXPIRY_SECONDS = 120 # Redis key expiry value
|
|
PING_EXPIRY_SECONDS = 5 # How long between PING and PONG before we raise the alarm - TODO - Impliment this
|
|
|
|
ASSIGN_INTERVAL_SECONDS = 30
|
|
LIVENESS_CHECK_INTERVAL_SECONDS = 5 # How often to cross check connected_workers with ping results from Redis
|
|
|
|
ENABLE_WEBSOCKET_PING_DEBUG = os.getenv('ENABLE_WEBSOCKET_PING_DEBUG', 'false').lower() == 'true'
|
|
|
|
ENABLE_TASK_ASSIGNMENT_DEBUG = os.getenv('ENABLE_TASK_ASSIGNMENT_DEBUG', 'false').lower() == 'true'
|
|
# Database configuration
|
|
DATABASE_URL = "mysql://root:password@172.17.0.1:3306/theapi"
|
|
|
|
# Connection pool for Redis
|
|
redis_connection_pool = redis.ConnectionPool.from_url(
|
|
REDIS_URL,
|
|
decode_responses=True,
|
|
max_connections=10000,
|
|
socket_timeout=5
|
|
)
|
|
|
|
def get_redis_client():
|
|
return redis.StrictRedis(connection_pool=redis_connection_pool)
|
|
|
|
# SQLAlchemy setup
|
|
pymysql.install_as_MySQLdb()
|
|
engine = create_engine(
|
|
DATABASE_URL,
|
|
isolation_level='SERIALIZABLE',
|
|
pool_size=5,
|
|
max_overflow=10,
|
|
pool_timeout=30,
|
|
pool_recycle=3600,
|
|
pool_pre_ping=True
|
|
)
|
|
|
|
SessionLocal = sessionmaker(
|
|
autocommit=False,
|
|
autoflush=False,
|
|
bind=engine,
|
|
expire_on_commit=False
|
|
)
|
|
|
|
Base = declarative_base()
|
|
|
|
# Logging setup
|
|
def setup_logging():
|
|
class FunctionNameFilter(logging.Filter):
|
|
def filter(self, record):
|
|
record.funcName = getattr(record, 'funcName', '<unknown>')
|
|
return True
|
|
|
|
log_format = (
|
|
"%(log_color)s%(asctime)s - %(levelname)s - %(funcName)s - %(message)s"
|
|
)
|
|
date_format = "%Y-%m-%d %H:%M:%S"
|
|
|
|
formatter = ColoredFormatter(
|
|
log_format,
|
|
datefmt=date_format,
|
|
log_colors={
|
|
"DEBUG": "cyan",
|
|
"INFO": "green",
|
|
"WARNING": "yellow",
|
|
"ERROR": "red",
|
|
"CRITICAL": "bold_red",
|
|
},
|
|
)
|
|
|
|
logger = logging.getLogger("websocket_server")
|
|
logger.setLevel(logging.DEBUG)
|
|
logger.addFilter(FunctionNameFilter())
|
|
|
|
stream_handler = StreamHandler()
|
|
stream_handler.setFormatter(formatter)
|
|
logger.addHandler(stream_handler)
|
|
|
|
# Ensure logs directory exists
|
|
log_dir = "logs"
|
|
os.makedirs(log_dir, exist_ok=True)
|
|
|
|
file_handler = FileHandler("logs/app.log")
|
|
file_handler.setFormatter(formatter)
|
|
file_handler.setLevel(logging.WARNING)
|
|
logger.addHandler(file_handler)
|
|
|
|
return logger
|
|
|
|
# Context manager for DB session
|
|
from contextlib import contextmanager
|
|
@contextmanager
|
|
def get_db_session():
|
|
session = SessionLocal()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|