From 2e205febf67207f5b132048d189741a493330cf9 Mon Sep 17 00:00:00 2001
From: Nguyen Viet Dung <29406816+magnified103@users.noreply.github.com>
Date: Wed, 25 Sep 2024 20:03:34 -0400
Subject: [PATCH] Report daily queue statistics to Discord (#409)

* Send daily queue statistics to Discord

* Timezone fix

* Lint fix
---
 additional_requirements.txt |  1 +
 dmoj/celery.py              | 12 +++++++
 dmoj/settings.py            |  3 ++
 judge/tasks/webhook.py      | 72 ++++++++++++++++++++++++++++++++++++-
 4 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/additional_requirements.txt b/additional_requirements.txt
index 4d6e9bb50..e6c99cc53 100644
--- a/additional_requirements.txt
+++ b/additional_requirements.txt
@@ -1,3 +1,4 @@
 uwsgi
 websocket-client
 watchdog
+matplotlib
diff --git a/dmoj/celery.py b/dmoj/celery.py
index e1da64064..d5c9a31a8 100644
--- a/dmoj/celery.py
+++ b/dmoj/celery.py
@@ -2,6 +2,7 @@
 import socket
 
 from celery import Celery
+from celery.schedules import crontab
 from celery.signals import task_failure
 
 app = Celery('dmoj')
@@ -20,6 +21,17 @@
 # Logger to enable reporting of errors.
 logger = logging.getLogger('judge.celery')
 
+# Load periodic tasks
+app.conf.beat_schedule = {
+    'daily-queue-time-stats': {
+        'task': 'judge.tasks.webhook.queue_time_stats',
+        'schedule': crontab(minute=0, hour=0),
+        'options': {
+            'expires': 60 * 60 * 24,
+        },
+    },
+}
+
 
 @task_failure.connect()
 def celery_failure_log(sender, task_id, exception, traceback, *args, **kwargs):
diff --git a/dmoj/settings.py b/dmoj/settings.py
index 0fb8e4121..14617a884 100755
--- a/dmoj/settings.py
+++ b/dmoj/settings.py
@@ -117,6 +117,8 @@
 
 VNOJ_LONG_QUEUE_ALERT_THRESHOLD = 10
 
+CELERY_TIMEZONE = 'Asia/Ho_Chi_Minh'
+
 # Some problems have a lot of testcases, and each testcase
 # has about 5~6 fields, so we need to raise this
 DATA_UPLOAD_MAX_NUMBER_FIELDS = 3000
@@ -180,6 +182,7 @@
     'on_new_blogpost': None,
     'on_error': None,
     'on_long_queue': None,
+    'queue_time_stats': None,
 }
 
 SITE_FULL_URL = None  # ie 'https://oj.vnoi.info', please remove the last / if needed
diff --git a/judge/tasks/webhook.py b/judge/tasks/webhook.py
index a4719306a..2eb3c41da 100644
--- a/judge/tasks/webhook.py
+++ b/judge/tasks/webhook.py
@@ -1,12 +1,18 @@
+import bisect
+from datetime import datetime, timedelta
+from io import BytesIO
+
 import pytz
 from celery import shared_task
 from discord_webhook import DiscordEmbed, DiscordWebhook
 from django.conf import settings
 from django.contrib.contenttypes.models import ContentType
+from django.db.models import F, FloatField
+from django.db.models.functions import Cast
 
 
 from judge.jinja2.gravatar import gravatar
-from judge.models import BlogPost, Comment, Contest, Problem, Tag, TagProblem, Ticket, TicketMessage
+from judge.models import BlogPost, Comment, Contest, Problem, Submission, Tag, TagProblem, Ticket, TicketMessage
 
 __all__ = ('on_new_ticket', 'on_new_comment', 'on_new_problem', 'on_new_tag_problem', 'on_new_tag', 'on_new_contest',
            'on_new_blogpost', 'on_new_ticket_message', 'on_long_queue')
@@ -208,3 +214,67 @@ def on_long_queue():
         return
 
     send_webhook(webhook, 'Long queue alert', None, None, url=f'{settings.SITE_FULL_URL}/submissions/?status=QU')
+
+
+@shared_task
+def queue_time_stats():
+    webhook_url = get_webhook_url('queue_time_stats')
+    if webhook_url is None:
+        return
+
+    end_time = (datetime.now(pytz.timezone(settings.CELERY_TIMEZONE))
+                .replace(hour=0, minute=0, second=0, microsecond=0))
+    start_time = end_time - timedelta(days=1)
+
+    queue_time = (Submission.objects.filter(date__gte=start_time, date__lte=end_time)
+                                    .filter(judged_date__isnull=False, rejudged_date__isnull=True)
+                                    .annotate(queue_time=Cast(F('judged_date') - F('date'), FloatField()) / 1000000.0)
+                                    .order_by('queue_time').values_list('queue_time', flat=True))
+
+    queue_time_ranges = [0, 1, 2, 5, 10, 30, 60, 120, 300, 600]
+    queue_time_labels = [
+        '',
+        '0s - 1s',
+        '1s - 2s',
+        '2s - 5s',
+        '5s - 10s',
+        '10s - 30s',
+        '30s - 1min',
+        '1min - 2min',
+        '2min - 5min',
+        '5min - 10min',
+        '> 10min',
+    ]
+
+    def binning(x):
+        return bisect.bisect_left(queue_time_ranges, x, lo=0, hi=len(queue_time_ranges))
+
+    queue_time_count = [0] * len(queue_time_labels)
+    for group in map(binning, list(queue_time)):
+        queue_time_count[group] += 1
+
+    import matplotlib
+    matplotlib.use('Agg')
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots()
+    y_pos = range(len(queue_time_labels) - 1)
+    bar = ax.barh(y_pos, queue_time_count[1:])
+    ax.bar_label(bar, labels=[x if x > 0 else '' for x in queue_time_count[1:]], padding=2)
+    ax.margins(x=0.2)
+    ax.set_yticks(y_pos, labels=queue_time_labels[1:])
+    ax.invert_yaxis()
+    ax.set_xlabel('Number of submissions')
+    fig.tight_layout()
+
+    with BytesIO() as f:
+        plt.savefig(f, format='png')
+        f.seek(0)
+        chart_bytes = f.read()
+
+    embed = DiscordEmbed(title=f'Queue time, {start_time:%Y-%m-%d}', color='03b2f8')
+    embed.set_image('attachment://chart.png')
+    webhook = DiscordWebhook(url=webhook_url)
+    webhook.add_file(chart_bytes, 'chart.png')
+    webhook.add_embed(embed)
+    webhook.execute()