From 132cf1da7f7c15a8cf4e183ee150c9f4b4d007d5 Mon Sep 17 00:00:00 2001
From: Ildar Iskhakov <Ildar.iskhakov@grafana.com>
Date: Tue, 29 Nov 2022 16:20:41 +0800
Subject: [PATCH] Add celery profiling (#913)

---
 engine/engine/celery.py | 25 +++++++++++++++++++++++--
 engine/requirements.txt |  1 +
 engine/settings/base.py |  2 ++
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/engine/engine/celery.py b/engine/engine/celery.py
index d21e968c..c023eb6e 100644
--- a/engine/engine/celery.py
+++ b/engine/engine/celery.py
@@ -1,13 +1,20 @@
+import logging
 import os
+import time
 
 import celery
 from celery.app.log import TaskFormatter
+from celery.utils.debug import memdump, sample_mem
+from celery.utils.log import get_task_logger
+from django.conf import settings
 
-# set the default Django settings module for the 'celery' program.
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings.prod")
 
 from django.db import connection  # noqa: E402
 
+logger = get_task_logger(__name__)
+logger.setLevel(logging.DEBUG)
+
 connection.cursor()
 from celery import Celery  # noqa: E402
 
@@ -38,6 +45,20 @@ def on_after_setup_logger(logger, **kwargs):
     for handler in logger.handlers:
         handler.setFormatter(
             TaskFormatter(
-                "%(asctime)s source=engine:celery task_id=%(task_id)s task_name=%(task_name)s name=%(name)s level=%(levelname)s %(message)s"
+                "%(asctime)s source=engine:celery worker=%(processName)s task_id=%(task_id)s task_name=%(task_name)s name=%(name)s level=%(levelname)s %(message)s"
             )
         )
+
+
+if settings.DEBUG_CELERY_TASKS_PROFILING:
+
+    @celery.signals.task_prerun.connect
+    def start_task_timer(task_id=None, task=None, *a, **kw):
+        logger.info("started: {} of {} with cpu={} at {}".format(task_id, task.name, time.perf_counter(), time.time()))
+        sample_mem()
+
+    @celery.signals.task_postrun.connect
+    def finish_task_timer(task_id=None, task=None, *a, **kw):
+        logger.info("ended: {} of {} with cpu={} at {}".format(task_id, task.name, time.perf_counter(), time.time()))
+        sample_mem()
+        memdump()
diff --git a/engine/requirements.txt b/engine/requirements.txt
index 3c5463f1..5e41b01d 100644
--- a/engine/requirements.txt
+++ b/engine/requirements.txt
@@ -41,3 +41,4 @@ psycopg2-binary==2.9.3
 emoji==1.7.0
 apns2==0.7.2
 regex==2021.11.2
+psutil==5.9.4
diff --git a/engine/settings/base.py b/engine/settings/base.py
index 1b1e9412..bbd17936 100644
--- a/engine/settings/base.py
+++ b/engine/settings/base.py
@@ -39,6 +39,8 @@ MIRAGE_CIPHER_MODE = "CBC"
 # SECURITY WARNING: don't run with debug turned on in production!
 DEBUG = False
 
+DEBUG_CELERY_TASKS_PROFILING = getenv_boolean("DEBUG_CELERY_TASKS_PROFILING", False)
+
 ALLOWED_HOSTS = [item.strip() for item in os.environ.get("ALLOWED_HOSTS", "*").split(",")]
 
 # TODO: update link to up-to-date docs