d4883256f1
Signed-off-by: zhanluxianshen <zhanluxianshen@163.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
import logging
|
|
from datetime import datetime
|
|
|
|
import click
|
|
from kombu.utils.url import parse_url # type: ignore
|
|
from redis import Redis
|
|
|
|
import app
|
|
from configs import dify_config
|
|
from extensions.ext_database import db
|
|
from libs.email_i18n import EmailType, get_email_i18n_service
|
|
|
|
redis_config = parse_url(dify_config.CELERY_BROKER_URL)
|
|
celery_redis = Redis(
|
|
host=redis_config.get("hostname") or "localhost",
|
|
port=redis_config.get("port") or 6379,
|
|
password=redis_config.get("password") or None,
|
|
db=int(redis_config.get("virtual_host")) if redis_config.get("virtual_host") else 1,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@app.celery.task(queue="monitor")
|
|
def queue_monitor_task():
|
|
queue_name = "dataset"
|
|
threshold = dify_config.QUEUE_MONITOR_THRESHOLD
|
|
|
|
if threshold is None:
|
|
logger.warning(click.style("QUEUE_MONITOR_THRESHOLD is not configured, skipping monitoring", fg="yellow"))
|
|
return
|
|
|
|
try:
|
|
queue_length = celery_redis.llen(f"{queue_name}")
|
|
logger.info(click.style(f"Start monitor {queue_name}", fg="green"))
|
|
|
|
if queue_length is None:
|
|
logger.error(
|
|
click.style(f"Failed to get queue length for {queue_name} - Redis may be unavailable", fg="red")
|
|
)
|
|
return
|
|
|
|
logger.info(click.style(f"Queue length: {queue_length}", fg="green"))
|
|
|
|
if queue_length >= threshold:
|
|
warning_msg = f"Queue {queue_name} task count exceeded the limit.: {queue_length}/{threshold}"
|
|
logging.warning(click.style(warning_msg, fg="red"))
|
|
alert_emails = dify_config.QUEUE_MONITOR_ALERT_EMAILS
|
|
if alert_emails:
|
|
to_list = alert_emails.split(",")
|
|
email_service = get_email_i18n_service()
|
|
for to in to_list:
|
|
try:
|
|
current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
email_service.send_email(
|
|
email_type=EmailType.QUEUE_MONITOR_ALERT,
|
|
language_code="en-US",
|
|
to=to,
|
|
template_context={
|
|
"queue_name": queue_name,
|
|
"queue_length": queue_length,
|
|
"threshold": threshold,
|
|
"alert_time": current_time,
|
|
},
|
|
)
|
|
except Exception as e:
|
|
logger.exception(click.style("Exception occurred during sending email", fg="red"))
|
|
|
|
except Exception as e:
|
|
logger.exception(click.style("Exception occurred during queue monitoring", fg="red"))
|
|
finally:
|
|
if db.session.is_active:
|
|
db.session.close()
|